[SCEVExpander] Make CanonicalMode handing in isSafeToExpand() more robust (PR50506)

[nikic]

Summary

isSafeToExpand() for addrecs depends on whether the SCEVExpander will be used in CanonicalMode. At least one caller currently gets this wrong, resulting in PR50506.

Fix this by a) making the CanonicalMode argument on the freestanding functions required and b) adding member functions on SCEVExpander that automatically take the SCEVExpander mode into account. We can use the latter variant nearly everywhere, and thus make sure that there is no chance of CanonicalMode mismatch.

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

[fhahn]

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

41

It looks like only 2 uses of those versions remain. I think it would be good to replace them with the member variants, to reduce the chance of mis-use. Although the use in LSR might require a bit of extra work :(

llvm/lib/Transforms/Scalar/LoopPredication.cpp

526–527

Could we just pass the used Expander and then use Expander.isSafeToExpandAt here?

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

3339

It might be worth creating the SCEVExpander up-front and then use the member version?

[reames]

LGTM to this patch as is. I agree with @fhahn's comments, but those are reasonable follow ups and don't need to be done in this patch. This is already is a significant improvement.

As a follow up thought, have you considered whether the member function variant of isSafeToExpand could be implemented in terms of isSafeToExpandAt with the internal insertion point? Not sure if that's a good idea or not...

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

 
 namespace llvm {
 extern cl::opt<unsigned> SCEVCheapExpansionBudget;
 
 /// Return true if the given expression is safe to expand in the sense that
 /// all materialized values are safe to speculate anywhere their operands are
 /// defined, and the expander is capable of expanding the expression.
 /// CanonicalMode indicates whether the expander will be used in canonical mode.
-bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE,
-                    bool CanonicalMode = true);
+bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE, bool CanonicalMode);
 
 /// Return true if the given expression is safe to expand in the sense that
 /// all materialized values are defined and safe to speculate at the specified
 /// location and their operands are defined at this location.
+/// CanonicalMode indicates whether the expander will be used in canonical mode.
 bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint,
-                      ScalarEvolution &SE);
+                      ScalarEvolution &SE, bool CanonicalMode);
 
 /// struct for holding enough information to help calculate the cost of the
 /// given SCEV when expanded into IR.
 struct SCEVOperand {
   explicit SCEVOperand(unsigned Opc, int Idx, const SCEV *S) :
     ParentOpcode(Opc), OperandIdx(Idx), S(S) { }
   /// LLVM instruction opcode that uses the operand.
   unsigned ParentOpcode;
#endif
   bool hoistIVInc(Instruction *IncV, Instruction *InsertPos);
 
   /// 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
+  /// all materialized values are safe to speculate anywhere their operands are
+  /// defined, and the expander is capable of expanding the expression.
+  bool isSafeToExpand(const SCEV *S) const;
+
+  /// Return true if the given expression is safe to expand in the sense that
+  /// all materialized values are defined and safe to speculate at the specified
+  /// location and their operands are defined at this location.
+  bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint) const;
+
   /// Insert code to directly compute the specified SCEV expression into the
   /// program.  The code is inserted into the specified block.
   Value *expandCodeFor(const SCEV *SH, Type *Ty, Instruction *I) {
     return expandCodeForImpl(SH, Ty, I, true);
   }
 
   /// Insert code to directly compute the specified SCEV expression into the
   /// program.  The code is inserted into the SCEVExpander's current

llvm/lib/CodeGen/HardwareLoops.cpp

if (SE.isLoopEntryGuardedByCond(L, ICmpInst::ICMP_NE, ExitCount,
     UseLoopGuard = false;
 
   BasicBlock *BB = L->getLoopPreheader();
   if (UseLoopGuard && BB->getSinglePredecessor() &&
       cast<BranchInst>(BB->getTerminator())->isUnconditional()) {
     BasicBlock *Predecessor = BB->getSinglePredecessor();
     // If it's not safe to create a while loop then don't force it and create a
     // do-while loop instead
-    if (!isSafeToExpandAt(ExitCount, Predecessor->getTerminator(), SE))
+    if (!SCEVE.isSafeToExpandAt(ExitCount, Predecessor->getTerminator()))
         UseLoopGuard = false;
     else
         BB = Predecessor;
   }
 
-  if (!isSafeToExpandAt(ExitCount, BB->getTerminator(), SE)) {
+  if (!SCEVE.isSafeToExpandAt(ExitCount, BB->getTerminator())) {
     LLVM_DEBUG(dbgs() << "- Bailing, unsafe to expand ExitCount "
                << *ExitCount << "\n");
     return nullptr;
   }
 
   Value *Count = SCEVE.expandCodeFor(ExitCount, CountType,
                                      BB->getTerminator());
 

llvm/lib/Target/PowerPC/PPCLoopInstrFormPrep.cpp

for (unsigned ChainIdx = 0; ChainIdx < Bucket.ChainBases.size(); ++ChainIdx) {
     unsigned BaseElemIdx = Bucket.ChainSize * ChainIdx;
     const SCEV *BaseSCEV =
         ChainIdx ? SE->getAddExpr(Bucket.BaseSCEV,
                                   Bucket.Elements[BaseElemIdx].Offset)
                  : Bucket.BaseSCEV;
     const SCEVAddRecExpr *BasePtrSCEV = cast<SCEVAddRecExpr>(BaseSCEV);
 
     // Make sure the base is able to expand.
-    if (!isSafeToExpand(BasePtrSCEV->getStart(), *SE))
+    if (!SCEVE.isSafeToExpand(BasePtrSCEV->getStart()))
       return MadeChange;
 
     assert(BasePtrSCEV->isAffine() &&
            "Invalid SCEV type for the base ptr for a candidate chain!\n");
 
     std::pair<Instruction *, Instruction *> Base = rewriteForBase(
         L, BasePtrSCEV, Bucket.Elements[BaseElemIdx].Instr,
         false /* CanPreInc */, ChainCommoning, SCEVE, DeletedPtrs);
for (unsigned Idx = BaseElemIdx + 1; Idx < BaseElemIdx + Bucket.ChainSize;
       const SCEV *OffsetSCEV =
           BaseElemIdx ? SE->getMinusSCEV(Bucket.Elements[Idx].Offset,
                                          Bucket.Elements[BaseElemIdx].Offset)
                       : Bucket.Elements[Idx].Offset;
 
       // Make sure offset is able to expand. Only need to check one time as the
       // offsets are reused between different chains.
       if (!BaseElemIdx)
-        if (!isSafeToExpand(OffsetSCEV, *SE))
+        if (!SCEVE.isSafeToExpand(OffsetSCEV))
           return false;
 
       Value *OffsetValue = SCEVE.expandCodeFor(
           OffsetSCEV, OffsetSCEV->getType(), LoopPredecessor->getTerminator());
 
       Instruction *NewPtr = rewriteForBucketElement(Base, Bucket.Elements[Idx],
                                                     OffsetValue, DeletedPtrs);
 
bool PPCLoopInstrFormPrep::rewriteLoadStores(
     PrepForm Form) {
   bool MadeChange = false;
 
   const SCEVAddRecExpr *BasePtrSCEV =
       cast<SCEVAddRecExpr>(BucketChain.BaseSCEV);
   if (!BasePtrSCEV->isAffine())
     return MadeChange;
 
-  if (!isSafeToExpand(BasePtrSCEV->getStart(), *SE))
-    return MadeChange;
-
-  SmallPtrSet<Value *, 16> DeletedPtrs;
-
   BasicBlock *Header = L->getHeader();
   SCEVExpander SCEVE(*SE, Header->getModule()->getDataLayout(),
                      "loopprepare-formrewrite");
+  if (!SCEVE.isSafeToExpand(BasePtrSCEV->getStart()))
+    return MadeChange;
+
+  SmallPtrSet<Value *, 16> DeletedPtrs;
 
   // For some DS form load/store instructions, it can also be an update form,
   // if the stride is constant and is a multipler of 4. Use update form if
   // prefer it.
   bool CanPreInc = (Form == UpdateForm ||
                     ((Form == DSForm) &&
                      isa<SCEVConstant>(BasePtrSCEV->getStepRecurrence(*SE)) &&
                      !cast<SCEVConstant>(BasePtrSCEV->getStepRecurrence(*SE))

llvm/lib/Transforms/Scalar/IndVarSimplify.cpp

bool IndVarSimplify::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
   // needing to kill the loop entirely.
   if (!LoopPredication)
     return false;
 
   // Note: ExactBTC is the exact backedge taken count *iff* the loop exits
   // through *explicit* control flow.  We have to eliminate the possibility of
   // implicit exits (see below) before we know it's truly exact.
   const SCEV *ExactBTC = SE->getBackedgeTakenCount(L);
-  if (isa<SCEVCouldNotCompute>(ExactBTC) || !isSafeToExpand(ExactBTC, *SE))
+  if (isa<SCEVCouldNotCompute>(ExactBTC) || !Rewriter.isSafeToExpand(ExactBTC))
     return false;
 
   assert(SE->isLoopInvariant(ExactBTC, L) && "BTC must be loop invariant");
   assert(ExactBTC->getType()->isIntegerTy() && "BTC must be integer");
 
   auto BadExit = [&](BasicBlock *ExitingBB) {
     // If our exiting block exits multiple loops, we can only rewrite the
     // innermost one.  Otherwise, we're changing how many times the innermost
auto BadExit = [&](BasicBlock *ExitingBB) {
     // within the loop which contains them.  This assumes trivially lcssa phis
     // have already been removed; TODO: generalize
     BasicBlock *ExitBlock =
     BI->getSuccessor(L->contains(BI->getSuccessor(0)) ? 1 : 0);
     if (!ExitBlock->phis().empty())
       return true;
 
     const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
-    if (isa<SCEVCouldNotCompute>(ExitCount) || !isSafeToExpand(ExitCount, *SE))
+    if (isa<SCEVCouldNotCompute>(ExitCount) ||
+        !Rewriter.isSafeToExpand(ExitCount))
       return true;
 
     assert(SE->isLoopInvariant(ExitCount, L) &&
            "Exit count must be loop invariant");
     assert(ExitCount->getType()->isIntegerTy() && "Exit count must be integer");
     return false;
   };
 

llvm/lib/Transforms/Scalar/InductiveRangeCheckElimination.cpp

else if (cannotBeMinInLoop(*SR.HighLimit, &OriginalLoop, SE,
       ExitPreLoopAtSCEV = SE.getAddExpr(*SR.HighLimit, MinusOneS);
     else {
       LLVM_DEBUG(dbgs() << "irce: could not prove no-overflow when computing "
                         << "preloop exit limit.  HighLimit = "
                         << *(*SR.HighLimit) << "\n");
       return false;
     }
 
-    if (!isSafeToExpandAt(ExitPreLoopAtSCEV, InsertPt, SE)) {
+    if (!Expander.isSafeToExpandAt(ExitPreLoopAtSCEV, InsertPt)) {
       LLVM_DEBUG(dbgs() << "irce: could not prove that it is safe to expand the"
                         << " preloop exit limit " << *ExitPreLoopAtSCEV
                         << " at block " << InsertPt->getParent()->getName()
                         << "\n");
       return false;
     }
 
     ExitPreLoopAt = Expander.expandCodeFor(ExitPreLoopAtSCEV, IVTy, InsertPt);
else if (cannotBeMinInLoop(*SR.LowLimit, &OriginalLoop, SE,
       ExitMainLoopAtSCEV = SE.getAddExpr(*SR.LowLimit, MinusOneS);
     else {
       LLVM_DEBUG(dbgs() << "irce: could not prove no-overflow when computing "
                         << "mainloop exit limit.  LowLimit = "
                         << *(*SR.LowLimit) << "\n");
       return false;
     }
 
-    if (!isSafeToExpandAt(ExitMainLoopAtSCEV, InsertPt, SE)) {
+    if (!Expander.isSafeToExpandAt(ExitMainLoopAtSCEV, InsertPt)) {
       LLVM_DEBUG(dbgs() << "irce: could not prove that it is safe to expand the"
                         << " main loop exit limit " << *ExitMainLoopAtSCEV
                         << " at block " << InsertPt->getParent()->getName()
                         << "\n");
       return false;
     }
 
     ExitMainLoopAt = Expander.expandCodeFor(ExitMainLoopAtSCEV, IVTy, InsertPt);

llvm/lib/Transforms/Scalar/LoopDataPrefetch.cpp

LLVM_DEBUG(dbgs() << "Loop has: "
              << (HasCall ? "calls" : "no calls") << ".\n");
 
   for (auto &P : Prefetches) {
     // Check if the stride of the accesses is large enough to warrant a
     // prefetch.
     if (!isStrideLargeEnough(P.LSCEVAddRec, TargetMinStride))
       continue;
 
+    BasicBlock *BB = P.InsertPt->getParent();
+    SCEVExpander SCEVE(*SE, BB->getModule()->getDataLayout(), "prefaddr");
     const SCEV *NextLSCEV = SE->getAddExpr(P.LSCEVAddRec, SE->getMulExpr(
       SE->getConstant(P.LSCEVAddRec->getType(), ItersAhead),
       P.LSCEVAddRec->getStepRecurrence(*SE)));
-    if (!isSafeToExpand(NextLSCEV, *SE))
+    if (!SCEVE.isSafeToExpand(NextLSCEV))
       continue;
 
-    BasicBlock *BB = P.InsertPt->getParent();
     Type *I8Ptr = Type::getInt8PtrTy(BB->getContext(), 0/*PtrAddrSpace*/);
-    SCEVExpander SCEVE(*SE, BB->getModule()->getDataLayout(), "prefaddr");
     Value *PrefPtrValue = SCEVE.expandCodeFor(NextLSCEV, I8Ptr, P.InsertPt);
 
     IRBuilder<> Builder(P.InsertPt);
     Module *M = BB->getParent()->getParent();
     Type *I32 = Type::getInt32Ty(BB->getContext());
     Function *PrefetchFunc = Intrinsic::getDeclaration(
         M, Intrinsic::prefetch, PrefPtrValue->getType());
     Builder.CreateCall(

llvm/lib/Transforms/Scalar/LoopIdiomRecognize.cpp

bool LoopIdiomRecognize::processLoopStridedStore(
   bool Changed = false;
   const SCEV *Start = Ev->getStart();
   // Handle negative strided loops.
   if (IsNegStride)
     Start = getStartForNegStride(Start, BECount, IntIdxTy, StoreSizeSCEV, SE);
 
   // TODO: ideally we should still be able to generate memset if SCEV expander
   // is taught to generate the dependencies at the latest point.
-  if (!isSafeToExpand(Start, *SE))
+  if (!Expander.isSafeToExpand(Start))
     return Changed;
 
   // Okay, we have a strided store "p[i]" of a splattable value.  We can turn
   // this into a memset in the loop preheader now if we want.  However, this
   // would be unsafe to do if there is anything else in the loop that may read
   // or write to the aliased location.  Check for any overlap by generating the
   // base pointer and checking the region.
   Value *BasePtr =
bool LoopIdiomRecognize::processLoopStridedStore(
 
   // Okay, everything looks good, insert the memset.
 
   const SCEV *NumBytesS =
       getNumBytes(BECount, IntIdxTy, StoreSizeSCEV, CurLoop, DL, SE);
 
   // TODO: ideally we should still be able to generate memset if SCEV expander
   // is taught to generate the dependencies at the latest point.
-  if (!isSafeToExpand(NumBytesS, *SE))
+  if (!Expander.isSafeToExpand(NumBytesS))
     return Changed;
 
   Value *NumBytes =
       Expander.expandCodeFor(NumBytesS, IntIdxTy, Preheader->getTerminator());
 
   CallInst *NewCall;
   if (SplatValue) {
     AAMDNodes AATags = TheStore->getAAMetadata();

llvm/lib/Transforms/Scalar/LoopPredication.cpp

 
 Instruction *LoopPredication::findInsertPt(Instruction *Use,
                                            ArrayRef<const SCEV*> Ops) {
   // Subtlety: SCEV considers things to be invariant if the value produced is
   // the same across iterations.  This is not the same as being able to
   // evaluate outside the loop, which is what we actually need here.
   for (const SCEV *Op : Ops)
     if (!SE->isLoopInvariant(Op, L) ||
-        !isSafeToExpandAt(Op, Preheader->getTerminator(), *SE))
+        !isSafeToExpandAt(Op, Preheader->getTerminator(), *SE,
+                          /* CanonicalMode */ true))
       return Use;
   return Preheader->getTerminator();
 }
 
 bool LoopPredication::isLoopInvariantValue(const SCEV* S) {
   // Handling expressions which produce invariant results, but *haven't* yet
   // been removed from the loop serves two important purposes.
   // 1) Most importantly, it resolves a pass ordering cycle which would
Optional<Value *> LoopPredication::widenICmpRangeCheckIncrementingLoop(
   // already guaranteed to dominate the guard.
   if (!isLoopInvariantValue(GuardStart) ||
       !isLoopInvariantValue(GuardLimit) ||
       !isLoopInvariantValue(LatchStart) ||
       !isLoopInvariantValue(LatchLimit)) {
     LLVM_DEBUG(dbgs() << "Can't expand limit check!\n");
     return None;
   }
-  if (!isSafeToExpandAt(LatchStart, Guard, *SE) ||
-      !isSafeToExpandAt(LatchLimit, Guard, *SE)) {
+  if (!Expander.isSafeToExpandAt(LatchStart, Guard) ||
+      !Expander.isSafeToExpandAt(LatchLimit, Guard)) {
     LLVM_DEBUG(dbgs() << "Can't expand limit check!\n");
     return None;
   }
 
   // guardLimit - guardStart + latchStart - 1
   const SCEV *RHS =
       SE->getAddExpr(SE->getMinusSCEV(GuardLimit, GuardStart),
                      SE->getMinusSCEV(LatchStart, SE->getOne(Ty)));
Optional<Value *> LoopPredication::widenICmpRangeCheckDecrementingLoop(
   // already guaranteed to dominate the guard.
   if (!isLoopInvariantValue(GuardStart) ||
       !isLoopInvariantValue(GuardLimit) ||
       !isLoopInvariantValue(LatchStart) ||
       !isLoopInvariantValue(LatchLimit)) {
     LLVM_DEBUG(dbgs() << "Can't expand limit check!\n");
     return None;
   }
-  if (!isSafeToExpandAt(LatchStart, Guard, *SE) ||
-      !isSafeToExpandAt(LatchLimit, Guard, *SE)) {
+  if (!Expander.isSafeToExpandAt(LatchStart, Guard) ||
+      !Expander.isSafeToExpandAt(LatchLimit, Guard)) {
     LLVM_DEBUG(dbgs() << "Can't expand limit check!\n");
     return None;
   }
   // The decrement of the latch check IV should be the same as the
   // rangeCheckIV.
   auto *PostDecLatchCheckIV = LatchCheck.IV->getPostIncExpr(*SE);
   if (RangeCheck.IV != PostDecLatchCheckIV) {
     LLVM_DEBUG(dbgs() << "Not the same. PostDecLatchCheckIV: "
bool LoopPredication::predicateLoopExits(Loop *L, SCEVExpander &Rewriter) {
   // important for profitability.  We may have a loop which hasn't been fully
   // canonicalized just yet.  If the exit we chose to widen is provably never
   // taken, we want the widened form to *also* be provably never taken.  We
   // can't guarantee this as a current unanalyzeable exit may later become
   // analyzeable, but we can at least avoid the obvious cases.
   const SCEV *MinEC = getMinAnalyzeableBackedgeTakenCount(*SE, *DT, L);
   if (isa<SCEVCouldNotCompute>(MinEC) || MinEC->getType()->isPointerTy() ||
       !SE->isLoopInvariant(MinEC, L) ||
-      !isSafeToExpandAt(MinEC, WidenableBR, *SE))
+      !Rewriter.isSafeToExpandAt(MinEC, WidenableBR))
     return ChangedLoop;
 
   // Subtlety: We need to avoid inserting additional uses of the WC.  We know
   // that it can only have one transitive use at the moment, and thus moving
   // that use to just before the branch and inserting code before it and then
   // modifying the operand is legal.
   auto *IP = cast<Instruction>(WidenableBR->getCondition());
   // Here we unconditionally modify the IR, so after this point we should return
for (BasicBlock *ExitingBB : ExitingBlocks) {
 
     // If already constant, nothing to do.
     if (isa<Constant>(BI->getCondition()))
       continue;
 
     const SCEV *ExitCount = SE->getExitCount(L, ExitingBB);
     if (isa<SCEVCouldNotCompute>(ExitCount) ||
         ExitCount->getType()->isPointerTy() ||
-        !isSafeToExpandAt(ExitCount, WidenableBR, *SE))
+        !Rewriter.isSafeToExpandAt(ExitCount, WidenableBR))
       continue;
 
     const bool ExitIfTrue = !L->contains(*succ_begin(ExitingBB));
     BasicBlock *ExitBB = BI->getSuccessor(ExitIfTrue ? 0 : 1);
     if (!ExitBB->getPostdominatingDeoptimizeCall())
       continue;
 
     /// Here we can be fairly sure that executing this exit will most likely

llvm/lib/Transforms/Scalar/LoopStrengthReduce.cpp

if (ICmpInst *CI = dyn_cast<ICmpInst>(UserInst)) {
           CI->setOperand(1, CI->getOperand(0));
           CI->setOperand(0, NV);
           NV = CI->getOperand(1);
           Changed = true;
         }
 
         // x == y  -->  x - y == 0
         const SCEV *N = SE.getSCEV(NV);
-        if (SE.isLoopInvariant(N, L) && isSafeToExpand(N, SE) &&
+        if (SE.isLoopInvariant(N, L) &&
+            isSafeToExpand(N, SE, /* CanonicalMode */ false) &&
             (!NV->getType()->isPointerTy() ||
              SE.getPointerBase(N) == SE.getPointerBase(S))) {
           // S is normalized, so normalize N before folding it into S
           // to keep the result normalized.
           N = normalizeForPostIncUse(N, TmpPostIncLoops, SE);
           Kind = LSRUse::ICmpZero;
           S = SE.getMinusSCEV(N, S);
         }

llvm/lib/Transforms/Utils/LoopUtils.cpp

while ((PN = dyn_cast<PHINode>(BBI++))) {
         // and varies predictably *inside* the loop.  Evaluate the value it
         // contains when the loop exits, if possible.  We prefer to start with
         // expressions which are true for all exits (so as to maximize
         // expression reuse by the SCEVExpander), but resort to per-exit
         // evaluation if that fails.
         const SCEV *ExitValue = SE->getSCEVAtScope(Inst, L->getParentLoop());
         if (isa<SCEVCouldNotCompute>(ExitValue) ||
             !SE->isLoopInvariant(ExitValue, L) ||
-            !isSafeToExpand(ExitValue, *SE)) {
+            !Rewriter.isSafeToExpand(ExitValue)) {
           // TODO: This should probably be sunk into SCEV in some way; maybe a
           // getSCEVForExit(SCEV*, L, ExitingBB)?  It can be generalized for
           // most SCEV expressions and other recurrence types (e.g. shift
           // recurrences).  Is there existing code we can reuse?
           const SCEV *ExitCount = SE->getExitCount(L, PN->getIncomingBlock(i));
           if (isa<SCEVCouldNotCompute>(ExitCount))
             continue;
           if (auto *AddRec = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(Inst)))
             if (AddRec->getLoop() == L)
               ExitValue = AddRec->evaluateAtIteration(ExitCount, *SE);
           if (isa<SCEVCouldNotCompute>(ExitValue) ||
               !SE->isLoopInvariant(ExitValue, L) ||
-              !isSafeToExpand(ExitValue, *SE))
+              !Rewriter.isSafeToExpand(ExitValue))
             continue;
         }
 
         // Computing the value outside of the loop brings no benefit if it is
         // definitely used inside the loop in a way which can not be optimized
         // away. Avoid doing so unless we know we have a value which computes
         // the ExitValue already. TODO: This should be merged into SCEV
         // expander to leverage its knowledge of existing expressions.

llvm/lib/Transforms/Utils/ScalarEvolutionExpander.cpp

for (PHINode *PN : PHIsToRemove) {
     InsertedValues.erase(PN);
     InsertedPostIncValues.erase(PN);
     PN->eraseFromParent();
   }
 
   return User->getOperand(OpIdx);
 }
 
+bool SCEVExpander::isSafeToExpand(const SCEV *S) const {
+  return llvm::isSafeToExpand(S, SE, CanonicalMode);
+}
+
+bool SCEVExpander::isSafeToExpandAt(const SCEV *S,
+                                    const Instruction *InsertionPoint) const {
+  return llvm::isSafeToExpandAt(S, InsertionPoint, SE, CanonicalMode);
+}
+
 namespace {
 // Search for a SCEV subexpression that is not safe to expand.  Any expression
 // that may expand to a !isSafeToSpeculativelyExecute value is unsafe, namely
 // UDiv expressions. We don't know if the UDiv is derived from an IR divide
 // instruction, but the important thing is that we prove the denominator is
 // nonzero before expansion.
 //
 // IVUsers already checks that IV-derived expressions are safe. So this check is

 namespace llvm {
 bool isSafeToExpand(const SCEV *S, ScalarEvolution &SE, bool CanonicalMode) {
   SCEVFindUnsafe Search(SE, CanonicalMode);
   visitAll(S, Search);
   return !Search.IsUnsafe;
 }
 
 bool isSafeToExpandAt(const SCEV *S, const Instruction *InsertionPoint,
-                      ScalarEvolution &SE) {
-  if (!isSafeToExpand(S, SE))
+                      ScalarEvolution &SE, bool CanonicalMode) {
+  if (!isSafeToExpand(S, SE, CanonicalMode))
     return false;
   // We have to prove that the expanded site of S dominates InsertionPoint.
   // This is easy when not in the same block, but hard when S is an instruction
   // to be expanded somewhere inside the same block as our insertion point.
   // What we really need here is something analogous to an OrderedBasicBlock,
   // but for the moment, we paper over the problem by handling two common and
   // cheap to check cases.
   if (SE.properlyDominates(S, InsertionPoint->getParent()))

llvm/lib/Transforms/Utils/SimplifyIndVar.cpp

if (!SE->isLoopInvariant(S, L))
     return false;
 
   // Do not generate something ridiculous even if S is loop invariant.
   if (Rewriter.isHighCostExpansion(S, L, SCEVCheapExpansionBudget, TTI, I))
     return false;
 
   auto *IP = GetLoopInvariantInsertPosition(L, I);
 
-  if (!isSafeToExpandAt(S, IP, *SE)) {
+  if (!Rewriter.isSafeToExpandAt(S, IP)) {
     LLVM_DEBUG(dbgs() << "INDVARS: Can not replace IV user: " << *I
                       << " with non-speculable loop invariant: " << *S << '\n');
     return false;
   }
 
   auto *Invariant = Rewriter.expandCodeFor(S, I->getType(), IP);
 
   I->replaceAllUsesWith(Invariant);

llvm/test/Transforms/IndVarSimplify/pr50506.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -S -indvars < %s | FileCheck %s
+
+; This test used to assert when expanding an addrec into a loop without
+; preheader.
+
+define void @test(ptr %tgt) {
+; CHECK-LABEL: @test(
+; CHECK-NEXT:  bb31:
+; CHECK-NEXT:    indirectbr ptr [[TGT:%.*]], [label [[EXIT:%.*]], label %bb33]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+; CHECK:       bb33:
+; CHECK-NEXT:    [[TMP34:%.*]] = phi i32 [ [[TMP50:%.*]], [[BB49:%.*]] ], [ 0, [[BB31:%.*]] ]
+; CHECK-NEXT:    br i1 false, label [[BB40_PREHEADER:%.*]], label [[BB49]]
+; CHECK:       bb40.preheader:
+; CHECK-NEXT:    br label [[BB40:%.*]]
+; CHECK:       bb40:
+; CHECK-NEXT:    br label [[BB40]]
+; CHECK:       bb49:
+; CHECK-NEXT:    [[TMP50]] = add i32 [[TMP34]], 1
+; CHECK-NEXT:    br label [[BB33:%.*]]
+;
+bb31:
+  indirectbr ptr %tgt, [label %exit, label %bb33]
+
+exit:
+  ret void
+
+bb33:                                             ; preds = %bb49, %bb31
+  %tmp34 = phi i32 [ %tmp50, %bb49 ], [ 0, %bb31 ]
+  %tmp36 = add i32 %tmp34, 1
+  br i1 false, label %bb40, label %bb49
+
+bb40:                                             ; preds = %bb38, %bb37
+  %tmp41 = phi i32 [ %tmp36, %bb33 ], [ %tmp39, %bb40 ]
+  %tmp39 = add i32 %tmp41, 0
+  br label %bb40
+
+bb49:                                             ; preds = %bb35
+  %tmp50 = add i32 %tmp34, 1
+  br label %bb33
+}

llvm/unittests/Transforms/Utils/ScalarEvolutionExpanderTest.cpp

TEST_F(ScalarEvolutionExpanderTest, SCEVExpanderIsSafeToExpandAt) {
   Builder.CreateCondBr(Cond, L, Post);
   Phi->addIncoming(ConstantInt::get(T_int64, 0), LPh);
   Phi->addIncoming(Add, L);
 
   Builder.SetInsertPoint(Post);
   Instruction *Ret = Builder.CreateRetVoid();
 
   ScalarEvolution SE = buildSE(*F);
+  SCEVExpander Exp(SE, M.getDataLayout(), "expander");
   const SCEV *S = SE.getSCEV(Phi);
   EXPECT_TRUE(isa<SCEVAddRecExpr>(S));
   const SCEVAddRecExpr *AR = cast<SCEVAddRecExpr>(S);
   EXPECT_TRUE(AR->isAffine());
-  EXPECT_FALSE(isSafeToExpandAt(AR, Top->getTerminator(), SE));
-  EXPECT_FALSE(isSafeToExpandAt(AR, LPh->getTerminator(), SE));
-  EXPECT_TRUE(isSafeToExpandAt(AR, L->getTerminator(), SE));
-  EXPECT_TRUE(isSafeToExpandAt(AR, Post->getTerminator(), SE));
+  EXPECT_FALSE(Exp.isSafeToExpandAt(AR, Top->getTerminator()));
+  EXPECT_FALSE(Exp.isSafeToExpandAt(AR, LPh->getTerminator()));
+  EXPECT_TRUE(Exp.isSafeToExpandAt(AR, L->getTerminator()));
+  EXPECT_TRUE(Exp.isSafeToExpandAt(AR, Post->getTerminator()));
 
   EXPECT_TRUE(LI->getLoopFor(L)->isLCSSAForm(*DT));
-  SCEVExpander Exp(SE, M.getDataLayout(), "expander");
   Exp.expandCodeFor(SE.getSCEV(Add), nullptr, Ret);
   EXPECT_TRUE(LI->getLoopFor(L)->isLCSSAForm(*DT));
 }
 
 // Check that SCEV expander does not use the nuw instruction
 // for expansion.
 TEST_F(ScalarEvolutionExpanderTest, SCEVExpanderNUW) {
   /*