[SCEV] Remove applyLoopGuards ExprsToRewrite

[caojoshua]

Summary

ExprsToRewrite is used to expand subexpressions, but it is unnecessary
and its purpose is not obvious. I changed the Guards Visitor so that it
can expand subexpressions in a single visit traversal. I think this
makes applyLoopGuards() much easier to read.

This change is not NFC. For urem guards, we directly overwrite
RewriteMap[LHS], but for other guards we write RewriteMap[RewrittenLHS].
Fix urem guards to do the same as other guards. In other words, expand
rewrites on top of existing rewrites, rather than completely overwrite
them.

I was not able to write a testcase for this change using
print<scalar-evolution> output. Manual dbgs() for this change verifies
that urem rewrites does not overwrite other rewrites.

For example, given:

%u = urem i32 %num, 4
%cmp = icmp eq i32 %u, 0
tail call void @llvm.assume(i1 %cmp)
%cmp.1 = icmp uge i32 %num, 4
tail call void @llvm.assume(i1 %cmp.1)
br label %for.body

The uge was getting overwritten by the urem. This change ensures both
assumes are applied.

[caojoshua]

llvm/lib/Analysis/ScalarEvolution.cpp

15035

NOTE: this is the only non-NFC change in this patch.

[caojoshua]

ping

[fhahn]

Is there a reason for why this was abandoned? Just curious if this is still a worthwhile simplification

[caojoshua]

In D142689#4416890, @fhahn wrote:

Is there a reason for why this was abandoned? Just curious if this is still a worthwhile simplification

The function got rewritten pretty dramatically and I did not want to rebase and rewrite at the time. I do think this work should still be possible.

llvm/lib/Analysis/ScalarEvolution.cpp

 
 /// A rewriter to replace SCEV expressions in Map with the corresponding entry
 /// in the map. It skips AddRecExpr because we cannot guarantee that the
 /// replacement is loop invariant in the loop of the AddRec.
 ///
 /// At the moment only rewriting SCEVUnknown and SCEVZeroExtendExpr is
 /// supported.
 class SCEVLoopGuardRewriter : public SCEVRewriteVisitor<SCEVLoopGuardRewriter> {
-  const DenseMap<const SCEV *, const SCEV *> &Map;
+  DenseMap<const SCEV *, const SCEV *> &Map;
 
 public:
   SCEVLoopGuardRewriter(ScalarEvolution &SE,
                         DenseMap<const SCEV *, const SCEV *> &M)
       : SCEVRewriteVisitor(SE), Map(M) {}
 
   const SCEV *visitAddRecExpr(const SCEVAddRecExpr *Expr) { return Expr; }
 
   const SCEV *visitUnknown(const SCEVUnknown *Expr) {
     auto I = Map.find(Expr);
     if (I == Map.end())
       return Expr;
     return I->second;
   }
 
   const SCEV *visitZeroExtendExpr(const SCEVZeroExtendExpr *Expr) {
     auto I = Map.find(Expr);
     if (I == Map.end())
       return SCEVRewriteVisitor<SCEVLoopGuardRewriter>::visitZeroExtendExpr(
           Expr);
-    return I->second;
+    Map.erase(Expr);
+    const SCEV *Result = SCEVRewriteVisitor<SCEVLoopGuardRewriter>::visit(I->second);
+    // It is possible that we visit Expr again and its result is cached. Erase
+    // it to ensure this visit's result is cached.
+    RewriteResults.erase(Expr);
+    return Result;
   }
 };
 
 const SCEV *ScalarEvolution::applyLoopGuards(const SCEV *Expr, const Loop *L) {
-  SmallVector<const SCEV *> ExprsToRewrite;
+  DenseMap<const SCEV *, const SCEV *> RewriteMap;
+
+  // Check whether LHS has already been rewritten. In that case we want to
+  // chain further rewrites onto the already rewritten value.
+  auto GetRewrittenLHS = [&](const SCEV *LHS) {
+    auto I = RewriteMap.find(LHS);
+    return I == RewriteMap.end() ? LHS : I->second;
+  };
+
   auto CollectCondition = [&](ICmpInst::Predicate Predicate, const SCEV *LHS,
-                              const SCEV *RHS,
-                              DenseMap<const SCEV *, const SCEV *>
-                                  &RewriteMap) {
+                              const SCEV *RHS) {
     // WARNING: It is generally unsound to apply any wrap flags to the proposed
     // replacement SCEV which isn't directly implied by the structure of that
     // SCEV.  In particular, using contextual facts to imply flags is *NOT*
     // legal.  See the scoping rules for flags in the header to understand why.
 
     // If LHS is a constant, apply information to the other expression.
     if (isa<SCEVConstant>(LHS)) {
       std::swap(LHS, RHS);
       Predicate = CmpInst::getSwappedPredicate(Predicate);
     }
 
     // Check for a condition of the form (-C1 + X < C2).  InstCombine will
     // create this form when combining two checks of the form (X u< C2 + C1) and
     // (X >=u C1).
     auto MatchRangeCheckIdiom = [this, Predicate, LHS, RHS, &RewriteMap,
-                                 &ExprsToRewrite]() {
+                                 &GetRewrittenLHS]() {
       auto *AddExpr = dyn_cast<SCEVAddExpr>(LHS);
       if (!AddExpr || AddExpr->getNumOperands() != 2)
         return false;
 
       auto *C1 = dyn_cast<SCEVConstant>(AddExpr->getOperand(0));
       auto *LHSUnknown = dyn_cast<SCEVUnknown>(AddExpr->getOperand(1));
       auto *C2 = dyn_cast<SCEVConstant>(RHS);
       if (!C1 || !C2 || !LHSUnknown)
         return false;
 
       auto ExactRegion =
           ConstantRange::makeExactICmpRegion(Predicate, C2->getAPInt())
               .sub(C1->getAPInt());
 
       // Bail out, unless we have a non-wrapping, monotonic range.
       if (ExactRegion.isWrappedSet() || ExactRegion.isFullSet())
         return false;
-      auto I = RewriteMap.find(LHSUnknown);
-      const SCEV *RewrittenLHS = I != RewriteMap.end() ? I->second : LHSUnknown;
+      const SCEV *RewrittenLHS = GetRewrittenLHS(LHSUnknown);
       RewriteMap[LHSUnknown] = getUMaxExpr(
           getConstant(ExactRegion.getUnsignedMin()),
           getUMinExpr(RewrittenLHS, getConstant(ExactRegion.getUnsignedMax())));
-      ExprsToRewrite.push_back(LHSUnknown);
       return true;
     };
     if (MatchRangeCheckIdiom())
       return;
 
     // If we have LHS == 0, check if LHS is computing a property of some unknown
     // SCEV %v which we can rewrite %v to express explicitly.
     const SCEVConstant *RHSC = dyn_cast<SCEVConstant>(RHS);
     if (Predicate == CmpInst::ICMP_EQ && RHSC &&
         RHSC->getValue()->isNullValue()) {
       // If LHS is A % B, i.e. A % B == 0, rewrite A to (A /u B) * B to
       // explicitly express that.
       const SCEV *URemLHS = nullptr;
       const SCEV *URemRHS = nullptr;
       if (matchURem(LHS, URemLHS, URemRHS)) {
         if (const SCEVUnknown *LHSUnknown = dyn_cast<SCEVUnknown>(URemLHS)) {
-          const auto *Multiple = getMulExpr(getUDivExpr(URemLHS, URemRHS), URemRHS);
+          const SCEV *RewrittenLHS = GetRewrittenLHS(LHSUnknown);
+          const auto *Multiple = getMulExpr(getUDivExpr(RewrittenLHS, URemRHS), URemRHS);
           RewriteMap[LHSUnknown] = Multiple;
-          ExprsToRewrite.push_back(LHSUnknown);
           return;
         }
       }
     }
 
     // Do not apply information for constants or if RHS contains an AddRec.
     if (isa<SCEVConstant>(LHS) || containsAddRecurrence(RHS))
       return;
 
     // If RHS is SCEVUnknown, make sure the information is applied to it.
     if (!isa<SCEVUnknown>(LHS) && isa<SCEVUnknown>(RHS)) {
       std::swap(LHS, RHS);
       Predicate = CmpInst::getSwappedPredicate(Predicate);
     }
 
     // Limit to expressions that can be rewritten.
     if (!isa<SCEVUnknown>(LHS) && !isa<SCEVZeroExtendExpr>(LHS))
       return;
 
-    // Check whether LHS has already been rewritten. In that case we want to
-    // chain further rewrites onto the already rewritten value.
-    auto I = RewriteMap.find(LHS);
-    const SCEV *RewrittenLHS = I != RewriteMap.end() ? I->second : LHS;
-
+    const SCEV *RewrittenLHS = GetRewrittenLHS(LHS);
     const SCEV *RewrittenRHS = nullptr;
     switch (Predicate) {
     case CmpInst::ICMP_ULT:
       RewrittenRHS =
           getUMinExpr(RewrittenLHS, getMinusSCEV(RHS, getOne(RHS->getType())));
       break;
     case CmpInst::ICMP_SLT:
       RewrittenRHS =
case CmpInst::ICMP_NE:
       if (isa<SCEVConstant>(RHS) &&
           cast<SCEVConstant>(RHS)->getValue()->isNullValue())
         RewrittenRHS = getUMaxExpr(RewrittenLHS, getOne(RHS->getType()));
       break;
     default:
       break;
     }
 
-    if (RewrittenRHS) {
+    if (RewrittenRHS)
       RewriteMap[LHS] = RewrittenRHS;
-      if (LHS == RewrittenLHS)
-        ExprsToRewrite.push_back(LHS);
-    }
   };
 
   BasicBlock *Header = L->getHeader();
   SmallVector<PointerIntPair<Value *, 1, bool>> Terms;
   // First, collect information from assumptions dominating the loop.
   for (auto &AssumeVH : AC.assumptions()) {
     if (!AssumeVH)
       continue;
for (std::pair<const BasicBlock *, const BasicBlock *> Pair(
     Terms.emplace_back(LoopEntryPredicate->getCondition(),
                        LoopEntryPredicate->getSuccessor(0) == Pair.second);
   }
 
   // Now apply the information from the collected conditions to RewriteMap.
   // Conditions are processed in reverse order, so the earliest conditions is
   // processed first. This ensures the SCEVs with the shortest dependency chains
   // are constructed first.
-  DenseMap<const SCEV *, const SCEV *> RewriteMap;
   for (auto [Term, EnterIfTrue] : reverse(Terms)) {
     SmallVector<Value *, 8> Worklist;
     SmallPtrSet<Value *, 8> Visited;
     Worklist.push_back(Term);
     while (!Worklist.empty()) {
       Value *Cond = Worklist.pop_back_val();
       if (!Visited.insert(Cond).second)
         continue;
 
       if (auto *Cmp = dyn_cast<ICmpInst>(Cond)) {
         auto Predicate =
             EnterIfTrue ? Cmp->getPredicate() : Cmp->getInversePredicate();
         const auto *LHS = getSCEV(Cmp->getOperand(0));
         const auto *RHS = getSCEV(Cmp->getOperand(1));
-        CollectCondition(Predicate, LHS, RHS, RewriteMap);
+        CollectCondition(Predicate, LHS, RHS);
         continue;
       }
 
       Value *L, *R;
       if (EnterIfTrue ? match(Cond, m_LogicalAnd(m_Value(L), m_Value(R)))
                       : match(Cond, m_LogicalOr(m_Value(L), m_Value(R)))) {
         Worklist.push_back(L);
         Worklist.push_back(R);
       }
     }
   }
 
   if (RewriteMap.empty())
     return Expr;
 
-  // Now that all rewrite information is collect, rewrite the collected
-  // expressions with the information in the map. This applies information to
-  // sub-expressions.
-  if (ExprsToRewrite.size() > 1) {
-    for (const SCEV *Expr : ExprsToRewrite) {
-      const SCEV *RewriteTo = RewriteMap[Expr];
-      RewriteMap.erase(Expr);
-      SCEVLoopGuardRewriter Rewriter(*this, RewriteMap);
-      RewriteMap.insert({Expr, Rewriter.visit(RewriteTo)});
-    }
-  }
-
   SCEVLoopGuardRewriter Rewriter(*this, RewriteMap);
   return Rewriter.visit(Expr);
 }