[IndVarSimplify] Wisely choose sext or zext when widening IV

[lihuang]

Summary

[IndVarSimplify] Wisely choose sext or zext when widening IV.

This is to fix the performance regression caused by two earlier patches: D18777 and D18867. (both were reverted).

The problem was, with those patches, some IV sexts are converted to zexts. IndVarSimplify marks the IV as signed or unsigned based on its cast users which comes first. If the IV's first cast user is a zext, it will be marked as unsigned. However, marking the IV as unsigned will cause the IV to be zero-extended during widening, if the IV also has sext users (e.g. a sext of IV for GEP), the sext cannot be removed and trunc is introduced. In general, we want to remove as many s/zexts as possible during widening.

The idea of this patch is to mark the IV as singed even if its first cast user is a zext, when

1. This IV is known to be non-negative
2. The first cast user, zext, has a GEP user. (GEP's offsets are treated as signed values)

[lihuang]

Sorry, didn't generate a full diff.

[reames]

I would suggest framing this slightly differently. If we know the IV is positive, we know that a sext and an zext must produce exactly the same result. We should canonicalize on one of them in general (which I think is what your previous patches caused) and the widening code should produce the optimal result for either. In fact, it looks like there's a NeverNegative flag on NarrowIVDefUse which already tries to do this.

I think the problem here is that we're failing to use this information in getWideRecurrence?

Note that I'm not an expert in this code and would really appreciate a second opinion.

[reames]

Another way of framing the possible fix is to change:

// Our raison d'etre! Eliminate sign and zero extension.
if (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse))

To:

if ((DU.NeverNegative && (isa<SExtInst>(DU.NarrowUse) || isa<ZExtInst>(DU.NarrowUse))) ||
   (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse)))

Obviously, clean up the code please. :)

This handles the sext/zext cases, but not the GEP cases. That still needs the getWideRecurrence parts.

[sanjoy]

Please add some test cases showing what specific cases you're trying to fix. Generally I agree with Philip's reasoning above -- this fix seems somewhat ad-hoc, if we _can_ prove that the narrow iv is positive, we should be able to take opportunistically pretend it was a signed or unsigned extension on a case by case basis (and get the best of both worlds).

[reames]

Also, as a separate patch, it would be good to add optimizations which cleanup the poor IR generated without the patch. In particular, we should be able to zext(trunc(sext(x))) and sext(trunc(zext(x)))) should be foldable when we know that x is a positive number within the range of the trunc output space. Both InstCombine (ValueTracking) and CVP (LVI) should be able to handle these cases. Same for the GEP cases.

[lihuang]

In D24280#535316, @reames wrote:

Another way of framing the possible fix is to change:

// Our raison d'etre! Eliminate sign and zero extension.
if (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse))

To:

if ((DU.NeverNegative && (isa<SExtInst>(DU.NarrowUse) || isa<ZExtInst>(DU.NarrowUse))) ||
   (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse)))

Obviously, clean up the code please. :)

This handles the sext/zext cases, but not the GEP cases. That still needs the getWideRecurrence parts.

This is the same logic as in D18867, but this is not enough, as you said, it doesn't work for the GEP cases.

[lihuang]

Hi Philip, Sanjoy,

Thank you for the suggestions. getWideRecurrence gives a zero-extended SCEV when IsSigned is false. We could also change it to return a sign-extended SCEV when value being extended is known non-negative.

I agree that we should canonicalize on one of the extends when the value being extended is non-negative.

[lihuang]

InstCombine canonicalizes on zexts when the value is non-negative. I will try canonicalizing on zext and teach getWideRecurrence to return a valid AddRec when the non-negative value is zero-extended.

[lihuang]

Canonicalize on zext when IV is known non-negative. Then, for each def-use pair of IV users, choose the kind of extension based on the following strategy:

if the narrow def is non-negative, use either sext or zext which applies to the narrow user.
otherwise, use the same extension applied to narrow def for the narrow user.

I added a map in class WidenIV to track the kind of extension used for each IV user, and basically changed every query to WideIV::IsSigned to a query into this map. The reason is, each narrow user should rely on its narrow def to choose sext/zext rather than the signedness of original IV, as the IV could be either sign extended or zero extended. Should be able to integrate this map with the Widened set, but I want to make sure I am in the right direction first.

This change includes the logic in D18867, so added those 2 test cases here. The 3rd test case comes from TSVC/174 (regression caused by D18777). All sext/zext to IV users in these 3 tests should be removed.

[lihuang]

Ping :),

Do you think this is the right direction?

Thanks,
-Li

[sanjoy]

Hi Li,

Sorry, I've been a bit busier than usual. I'll take a look at this this weekend.

• Sanjoy

[lihuang]

Hi Sanjoy,

No problem, take your time.

[sanjoy]

Some comments inline.

lib/Transforms/Scalar/IndVarSimplify.cpp

916

Please name the name so that the intent of the bool is obvious (perhaps WasSignExtended?). Or (even better) use something other than a bool as the domain (perhaps an enum?). Also, clang-format.

946

I'm not a big fan of returning via params like these -- it is too easy to forget to update the return value along one path. Have you considered returning an std::pair or even a hand-rolled struct instead?

954

Document these please.

1163

Common these three ExtendKindMap[DU.NarrowDef] lookups into one with a descriptive name.

Actually, these should not look up the DenseMap directly like these -- if DU.NarrowDef is not in ExtendKindMap then we'll get a bogus false result here. You should add an accessor that will trip an assert if DU.NarrowDef is not present in ExtendKindMap.

1197

Here and elsewhere please document the new WidenBySext paramenter.

1220

clang-format please

1391

I'd spell it as canWidenBySExt (same below).

1470

Does ExtendKindMap subsume the IsSigned field? If so, that field should be removed.

[lihuang]

Address Sanjoy's comments.

[lihuang]

Add another test case where zext should not be removed.

lib/Transforms/Scalar/IndVarSimplify.cpp

916

Sorry for the late reply, I was working on other stuff.

Added an enum type ExtendKind.

946

thanks for the suggestion, I changed the return type to a std::pair of AddRec* and ExtendKind.

1163–1164

Actually DU.NarrowDef should be already in this map when we come to visit NarrowUse, because we visit NarrowDef before pushing NarrowUse to the queue when doing BFS.

I added an assert in WidenIVUse to ensure this.

1194–1198

Removed this parameter.

1470

Yes it does, removed the IsSigned field from WidenIV.

[sanjoy]

Hi Li,

This is looking fairly close to ready. I mostly have some nitpicky comments inline.

lib/Transforms/Scalar/IndVarSimplify.cpp

919

Make the domain of this map an AssertingVH to guard against use-after-free bugs.

1163–1164

Actually DU.NarrowDef should be already in this map when we come to visit NarrowUse, because we visit NarrowDef before pushing NarrowUse to the queue when doing BFS.

Yes, but I'd rather have the assert automatically trigger anything you access ExtendKindMap.

1202

s/SignExtended/Unknown/

Not compulsory to change: you can use {A, B} instead of std::make_pair(A, B).

1306

Have you considered putting canWidenBySExt and canWidenByZExt as lambdas here? That will make things more readable I think.

1347

Add an assert((WideAddRec.first == nullptr) == (.second == Unknown)).

[lihuang]

lib/Transforms/Scalar/IndVarSimplify.cpp

1306–1307

Good suggestion :) I changed them to lambdas

[sanjoy]

Only one thing left. :)

lib/Transforms/Scalar/IndVarSimplify.cpp

1163–1164

This was not done.

To be clear: I was suggesting adding a method:

ExtendKind getExtendKind(Instruction *I) {
  auto It = ExtendKindMap.find(I);
  assert(It != ExtendKindMap.end() && "Not found!");
  return It->second;
}

to the WidenIV class, and using that instead of ExtendKindMap[DU.NarrowDef] here and below.

[lihuang]

lib/Transforms/Scalar/IndVarSimplify.cpp

1163–1164

I see...

I think DenseMap should have an "at" like std::map::at, but it doesn't. Okay, I will add a method like you are suggesting here.

[sanjoy]

Lgtm

[evstupac]

Committed.
Revision 282650

lib/Transforms/Scalar/IndVarSimplify.cpp

 /// Update information about the induction variable that is extended by this
 /// sign or zero extend operation. This is used to determine the final width of
 /// the IV before actually widening it.
 static void visitIVCast(CastInst *Cast, WideIVInfo &WI, ScalarEvolution *SE,
                         const TargetTransformInfo *TTI) {
   bool IsSigned = Cast->getOpcode() == Instruction::SExt;
   if (!IsSigned && Cast->getOpcode() != Instruction::ZExt)
     return;
   
+  // If this is a zext, mark this IV as signed if it is known to be
+  // non-negative and the zext has a GEP user.
+  if (!IsSigned) {
+    const SCEV *NarrowIVSCEV = SE->getSCEV(WI.NarrowIV);
+    bool NeverNegative = 
+      SE->isKnownPredicate(ICmpInst::ICMP_SGE, NarrowIVSCEV,
+                           SE->getConstant(NarrowIVSCEV->getType(), 0));
+    if (NeverNegative && any_of(Cast->users(), [](User *U) {
+      auto* CastUser = cast<Instruction>(U);
+      return CastUser->getOpcode() == Instruction::GetElementPtr;
+    }))
+      IsSigned = true;
+  }
+
   Type *Ty = Cast->getType();
   uint64_t Width = SE->getTypeSizeInBits(Ty);
   if (!Cast->getModule()->getDataLayout().isLegalInteger(Width))
     return;
 
   // Check that `Cast` actually extends the induction variable (we rely on this
   // later).  This takes care of cases where `Cast` is extending a truncation of
   // the narrow induction variable, and thus can end up being narrower than the
class WidenIV {
   const SCEV *WideIncExpr;
   SmallVectorImpl<WeakVH> &DeadInsts;
 
   SmallPtrSet<Instruction*,16> Widened;
   SmallVector<NarrowIVDefUse, 8> NarrowIVUsers;
 
 public:
   WidenIV(const WideIVInfo &WI, LoopInfo *LInfo,
           ScalarEvolution *SEv, DominatorTree *DTree,
           SmallVectorImpl<WeakVH> &DI) :
     OrigPhi(WI.NarrowIV),
     WideType(WI.WidestNativeType),
     IsSigned(WI.IsSigned),
     LI(LInfo),
     L(LI->getLoopFor(OrigPhi->getParent())),
     SE(SEv),
     DT(DTree),
     WidePhi(nullptr),
     WideInc(nullptr),
     WideIncExpr(nullptr),
protected:
   Instruction *cloneIVUser(NarrowIVDefUse DU, const SCEVAddRecExpr *WideAR);
   Instruction *cloneArithmeticIVUser(NarrowIVDefUse DU,
                                      const SCEVAddRecExpr *WideAR);
   Instruction *cloneBitwiseIVUser(NarrowIVDefUse DU);
 
   const SCEVAddRecExpr *getWideRecurrence(Instruction *NarrowUse);
 
   const SCEVAddRecExpr* getExtendedOperandRecurrence(NarrowIVDefUse DU);
 
   const SCEV *getSCEVByOpCode(const SCEV *LHS, const SCEV *RHS,
                               unsigned OpCode) const;
 
   Instruction *widenIVUse(NarrowIVDefUse DU, SCEVExpander &Rewriter);
 
   bool widenLoopCompare(NarrowIVDefUse DU);
 
   void pushNarrowIVUsers(Instruction *NarrowDef, Instruction *WideDef);
 };
 } // anonymous namespace
 
 /// Perform a quick domtree based check for loop invariance assuming that V is
 /// used within the loop. LoopInfo::isLoopInvariant() seems gratuitous for this
 /// purpose.
 static bool isLoopInvariant(Value *V, const Loop *L, const DominatorTree *DT) {
   Instruction *Inst = dyn_cast<Instruction>(V);
const SCEVAddRecExpr* WidenIV::getExtendedOperandRecurrence(NarrowIVDefUse DU) {
   // if extending the other will lead to a recurrence.
   const unsigned ExtendOperIdx =
       DU.NarrowUse->getOperand(0) == DU.NarrowDef ? 1 : 0;
   assert(DU.NarrowUse->getOperand(1-ExtendOperIdx) == DU.NarrowDef && "bad DU");
 
   const SCEV *ExtendOperExpr = nullptr;
   const OverflowingBinaryOperator *OBO =
     cast<OverflowingBinaryOperator>(DU.NarrowUse);
   if (IsSigned && OBO->hasNoSignedWrap())
     ExtendOperExpr = SE->getSignExtendExpr(
       SE->getSCEV(DU.NarrowUse->getOperand(ExtendOperIdx)), WideType);
   else if(!IsSigned && OBO->hasNoUnsignedWrap())
     ExtendOperExpr = SE->getZeroExtendExpr(
       SE->getSCEV(DU.NarrowUse->getOperand(ExtendOperIdx)), WideType);
   else
     return nullptr;
 
   // When creating this SCEV expr, don't apply the current operations NSW or NUW
const SCEVAddRecExpr* WidenIV::getExtendedOperandRecurrence(NarrowIVDefUse DU) {
 
   if (!AddRec || AddRec->getLoop() != L)
     return nullptr;
   return AddRec;
 }
 
 /// Is this instruction potentially interesting for further simplification after
 /// widening it's type? In other words, can the extend be safely hoisted out of
 /// the loop with SCEV reducing the value to a recurrence on the same loop. If
 /// so, return the sign or zero extended recurrence. Otherwise return NULL.
 const SCEVAddRecExpr *WidenIV::getWideRecurrence(Instruction *NarrowUse) {
   if (!SE->isSCEVable(NarrowUse->getType()))
     return nullptr;
 
   const SCEV *NarrowExpr = SE->getSCEV(NarrowUse);
   if (SE->getTypeSizeInBits(NarrowExpr->getType()) >=
       SE->getTypeSizeInBits(WideType)) {
     // NarrowUse implicitly widens its operand. e.g. a gep with a narrow
     // index. So don't follow this use.
     return nullptr;
   }
 
   const SCEV *WideExpr = IsSigned ?
     SE->getSignExtendExpr(NarrowExpr, WideType) :
     SE->getZeroExtendExpr(NarrowExpr, WideType);
   const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(WideExpr);
   if (!AddRec || AddRec->getLoop() != L)
     return nullptr;
   return AddRec;
 }
 
 /// This IV user cannot be widen. Replace this use of the original narrow IV
 /// with a truncation of the new wide IV to isolate and eliminate the narrow IV.
 static void truncateIVUse(NarrowIVDefUse DU, DominatorTree *DT, LoopInfo *LI) {
   DEBUG(dbgs() << "INDVARS: Truncate IV " << *DU.WideDef
         << " for user " << *DU.NarrowUse << "\n");
   IRBuilder<> Builder(
       getInsertPointForUses(DU.NarrowUse, DU.NarrowDef, DT, LI));
   Value *Trunc = Builder.CreateTrunc(DU.WideDef, DU.NarrowDef->getType());
   DU.NarrowUse->replaceUsesOfWith(DU.NarrowDef, Trunc);
 }
 
 /// If the narrow use is a compare instruction, then widen the compare
if (LI->getLoopFor(UsePhi->getParent()) != L) {
         DeadInsts.emplace_back(UsePhi);
         DEBUG(dbgs() << "INDVARS: Widen lcssa phi " << *UsePhi
               << " to " << *WidePhi << "\n");
       }
       return nullptr;
     }
   }
   // Our raison d'etre! Eliminate sign and zero extension.
   if (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse)) {
     Value *NewDef = DU.WideDef;
     if (DU.NarrowUse->getType() != WideType) {
       unsigned CastWidth = SE->getTypeSizeInBits(DU.NarrowUse->getType());
       unsigned IVWidth = SE->getTypeSizeInBits(WideType);
       if (CastWidth < IVWidth) {
         // The cast isn't as wide as the IV, so insert a Trunc.
         IRBuilder<> Builder(DU.NarrowUse);
         NewDef = Builder.CreateTrunc(DU.WideDef, DU.NarrowUse->getType());
       }
if (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse)) {
     // No further widening is needed. The deceased [sz]ext had done it for us.
     return nullptr;
   }
 
   // Does this user itself evaluate to a recurrence after widening?
   const SCEVAddRecExpr *WideAddRec = getWideRecurrence(DU.NarrowUse);
   if (!WideAddRec)
     WideAddRec = getExtendedOperandRecurrence(DU);
 
   if (!WideAddRec) {
     // If use is a loop condition, try to promote the condition instead of
     // truncating the IV first.
     if (widenLoopCompare(DU))
       return nullptr;
 
     // This user does not evaluate to a recurence after widening, so don't
     // follow it. Instead insert a Trunc to kill off the original use,
if (WideAddRec != SE->getSCEV(WideUse)) {
     DeadInsts.emplace_back(WideUse);
     return nullptr;
   }
 
   // Returning WideUse pushes it on the worklist.
   return WideUse;
 }
 
 /// Add eligible users of NarrowDef to NarrowIVUsers.
 ///
 void WidenIV::pushNarrowIVUsers(Instruction *NarrowDef, Instruction *WideDef) {
   const SCEV *NarrowSCEV = SE->getSCEV(NarrowDef);
   bool NeverNegative =
       SE->isKnownPredicate(ICmpInst::ICMP_SGE, NarrowSCEV,
                            SE->getConstant(NarrowSCEV->getType(), 0));
   for (User *U : NarrowDef->users()) {
     Instruction *NarrowUser = cast<Instruction>(U);
PHINode *WidenIV::createWideIV(SCEVExpander &Rewriter) {
   DEBUG(dbgs() << "Wide IV: " << *WidePhi << "\n");
   ++NumWidened;
 
   // Traverse the def-use chain using a worklist starting at the original IV.
   assert(Widened.empty() && NarrowIVUsers.empty() && "expect initial state" );
 
   Widened.insert(OrigPhi);
   pushNarrowIVUsers(OrigPhi, WidePhi);
 
   while (!NarrowIVUsers.empty()) {
     NarrowIVDefUse DU = NarrowIVUsers.pop_back_val();
 
     // Process a def-use edge. This may replace the use, so don't hold a
     // use_iterator across it.
     Instruction *WideUse = widenIVUse(DU, Rewriter);
 
     // Follow all def-use edges from the previous narrow use.