[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

901

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.

940

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?

948

Document these please.

1159–1160

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.

1192–1193

Here and elsewhere please document the new WidenBySext paramenter.

1217

clang-format please

1420

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

1499

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

901

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

Added an enum type ExtendKind.

940

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

1159–1161

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.

1192–1197

Removed this parameter.

1499

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

904

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

1159–1161

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.

1197

s/SignExtended/Unknown/

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

1331–1332

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

1373

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

[lihuang]

lib/Transforms/Scalar/IndVarSimplify.cpp

1331–1332

Good suggestion :) I changed them to lambdas

[sanjoy]

Only one thing left. :)

lib/Transforms/Scalar/IndVarSimplify.cpp

1159–1161

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

1159–1161

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

 /// creating any new induction variables. To do this, it creates a new phi of
 /// the wider type and redirects all users, either removing extends or inserting
 /// truncs whenever we stop propagating the type.
 ///
 class WidenIV {
   // Parameters
   PHINode *OrigPhi;
   Type *WideType;
-  bool IsSigned;
 
   // Context
   LoopInfo        *LI;
   Loop            *L;
   ScalarEvolution *SE;
   DominatorTree   *DT;
 
   // Result
   PHINode *WidePhi;
   Instruction *WideInc;
   const SCEV *WideIncExpr;
   SmallVectorImpl<WeakVH> &DeadInsts;
 
-  SmallPtrSet<Instruction*,16> Widened;
+  SmallPtrSet<Instruction *,16> Widened;
   SmallVector<NarrowIVDefUse, 8> NarrowIVUsers;
 
+  enum ExtendKind { ZeroExtended, SignExtended, Unknown };
+  // A map tracking the kind of extension used to widen each narrow IV 
+  // and narrow IV user.
+  // Key: pointer to a narrow IV or IV user.
+  // Value: the kind of extension used to widen this Instruction.
+  DenseMap<AssertingVH<Instruction>, ExtendKind> ExtendKindMap;
+
 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),
     DeadInsts(DI) {
     assert(L->getHeader() == OrigPhi->getParent() && "Phi must be an IV");
+    ExtendKindMap[OrigPhi] = WI.IsSigned ? SignExtended : ZeroExtended;
   }
 
   PHINode *createWideIV(SCEVExpander &Rewriter);
 
 protected:
   Value *createExtendInst(Value *NarrowOper, Type *WideType, bool IsSigned,
                           Instruction *Use);
 
   Instruction *cloneIVUser(NarrowIVDefUse DU, const SCEVAddRecExpr *WideAR);
   Instruction *cloneArithmeticIVUser(NarrowIVDefUse DU,
                                      const SCEVAddRecExpr *WideAR);
   Instruction *cloneBitwiseIVUser(NarrowIVDefUse DU);
 
-  const SCEVAddRecExpr *getWideRecurrence(Instruction *NarrowUse);
+  typedef std::pair<const SCEVAddRecExpr *, ExtendKind> WidenedRecTy;
+
+  WidenedRecTy getWideRecurrence(NarrowIVDefUse DU);
 
-  const SCEVAddRecExpr* getExtendedOperandRecurrence(NarrowIVDefUse DU);
+  WidenedRecTy 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);
Instruction *WidenIV::cloneBitwiseIVUser(NarrowIVDefUse DU) {
   Instruction *WideDef = DU.WideDef;
 
   DEBUG(dbgs() << "Cloning bitwise IVUser: " << *NarrowUse << "\n");
 
   // Replace NarrowDef operands with WideDef. Otherwise, we don't know anything
   // about the narrow operand yet so must insert a [sz]ext. It is probably loop
   // invariant and will be folded or hoisted. If it actually comes from a
   // widened IV, it should be removed during a future call to widenIVUse.
+  bool IsSigned = ExtendKindMap[NarrowDef] == SignExtended;
   Value *LHS = (NarrowUse->getOperand(0) == NarrowDef)
                    ? WideDef
                    : createExtendInst(NarrowUse->getOperand(0), WideType,
                                       IsSigned, NarrowUse);
   Value *RHS = (NarrowUse->getOperand(1) == NarrowDef)
                    ? WideDef
                    : createExtendInst(NarrowUse->getOperand(1), WideType,
                                       IsSigned, NarrowUse);
auto GuessNonIVOperand = [&](bool SignExt) {
     case Instruction::Sub:
       WideUse = SE->getMinusSCEV(WideLHS, WideRHS);
       break;
     }
 
     return WideUse == WideAR;
   };
 
-  bool SignExtend = IsSigned;
+  bool SignExtend = ExtendKindMap[NarrowDef] == SignExtended;
   if (!GuessNonIVOperand(SignExtend)) {
     SignExtend = !SignExtend;
     if (!GuessNonIVOperand(SignExtend))
       return nullptr;
   }
 
   Value *LHS = (NarrowUse->getOperand(0) == NarrowDef)
                    ? WideDef
if (OpCode == Instruction::Mul)
     return SE->getMulExpr(LHS, RHS);
 
   llvm_unreachable("Unsupported opcode.");
 }
 
 /// No-wrap operations can transfer sign extension of their result to their
 /// operands. Generate the SCEV value for the widened operation without
 /// actually modifying the IR yet. If the expression after extending the
-/// operands is an AddRec for this loop, return it.
-const SCEVAddRecExpr* WidenIV::getExtendedOperandRecurrence(NarrowIVDefUse DU) {
+/// operands is an AddRec for this loop, return the AddRec and the kind of
+/// extension used.
+WidenIV::WidenedRecTy WidenIV::getExtendedOperandRecurrence(NarrowIVDefUse DU) {
 
   // Handle the common case of add<nsw/nuw>
   const unsigned OpCode = DU.NarrowUse->getOpcode();
   // Only Add/Sub/Mul instructions supported yet.
   if (OpCode != Instruction::Add && OpCode != Instruction::Sub &&
       OpCode != Instruction::Mul)
-    return nullptr;
+    return {nullptr, Unknown};
 
   // One operand (NarrowDef) has already been extended to WideDef. Now determine
   // 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())
+  ExtendKind ExtKind = ExtendKindMap[DU.NarrowDef];
+  if (ExtKind == SignExtended && OBO->hasNoSignedWrap())
     ExtendOperExpr = SE->getSignExtendExpr(
       SE->getSCEV(DU.NarrowUse->getOperand(ExtendOperIdx)), WideType);
-  else if(!IsSigned && OBO->hasNoUnsignedWrap())
+  else if(ExtKind == ZeroExtended && OBO->hasNoUnsignedWrap())
     ExtendOperExpr = SE->getZeroExtendExpr(
       SE->getSCEV(DU.NarrowUse->getOperand(ExtendOperIdx)), WideType);
   else
-    return nullptr;
+    return {nullptr, Unknown};
 
   // When creating this SCEV expr, don't apply the current operations NSW or NUW
   // flags. This instruction may be guarded by control flow that the no-wrap
   // behavior depends on. Non-control-equivalent instructions can be mapped to
   // the same SCEV expression, and it would be incorrect to transfer NSW/NUW
   // semantics to those operations.
   const SCEV *lhs = SE->getSCEV(DU.WideDef);
   const SCEV *rhs = ExtendOperExpr;
 
   // Let's swap operands to the initial order for the case of non-commutative
   // operations, like SUB. See PR21014.
   if (ExtendOperIdx == 0)
     std::swap(lhs, rhs);
   const SCEVAddRecExpr *AddRec =
       dyn_cast<SCEVAddRecExpr>(getSCEVByOpCode(lhs, rhs, OpCode));
 
   if (!AddRec || AddRec->getLoop() != L)
-    return nullptr;
-  return AddRec;
+    return {nullptr, Unknown};
+
+  return {AddRec, ExtKind};
 }
 
 /// 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;
+/// so, return the extended recurrence and the kind of extension used. Otherwise
+/// return {nullptr, Unknown}.
+WidenIV::WidenedRecTy WidenIV::getWideRecurrence(NarrowIVDefUse DU) {
+  if (!SE->isSCEVable(DU.NarrowUse->getType()))
+    return {nullptr, Unknown};
 
-  const SCEV *NarrowExpr = SE->getSCEV(NarrowUse);
+  const SCEV *NarrowExpr = SE->getSCEV(DU.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;
+    return {nullptr, Unknown};
   }
 
-  const SCEV *WideExpr = IsSigned ?
-    SE->getSignExtendExpr(NarrowExpr, WideType) :
-    SE->getZeroExtendExpr(NarrowExpr, WideType);
+  const SCEV *WideExpr;
+  ExtendKind ExtKind;
+  if (DU.NeverNegative) {
+    WideExpr = SE->getSignExtendExpr(NarrowExpr, WideType);
+    if (isa<SCEVAddRecExpr>(WideExpr))
+      ExtKind = SignExtended;
+    else {
+      WideExpr = SE->getZeroExtendExpr(NarrowExpr, WideType);
+      ExtKind = ZeroExtended;
+    }
+  } else if (ExtendKindMap[DU.NarrowDef] == SignExtended) {
+    WideExpr = SE->getSignExtendExpr(NarrowExpr, WideType);
+    ExtKind = SignExtended;
+  } else {
+    WideExpr = SE->getZeroExtendExpr(NarrowExpr, WideType);
+    ExtKind = ZeroExtended;
+  }
   const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(WideExpr);
   if (!AddRec || AddRec->getLoop() != L)
-    return nullptr;
-  return AddRec;
+    return {nullptr, Unknown};
+  return {AddRec, ExtKind};
 }
 
 /// 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(
bool WidenIV::widenLoopCompare(NarrowIVDefUse DU) {
   //    %narrow for the following use
   //
   //      icmp slt i32 %narrow, %val   ... (A)
   //
   //    and %narrow is always positive.  Then
   //
   //      (A) == icmp slt i32 sext(%narrow), sext(%val)
   //          == icmp slt i32 zext(%narrow), sext(%val)
-
+  bool IsSigned = ExtendKindMap[DU.NarrowDef] == SignExtended;
   if (!(DU.NeverNegative || IsSigned == Cmp->isSigned()))
     return false;
 
   Value *Op = Cmp->getOperand(Cmp->getOperand(0) == DU.NarrowDef ? 1 : 0);
   unsigned CastWidth = SE->getTypeSizeInBits(Op->getType());
   unsigned IVWidth = SE->getTypeSizeInBits(WideType);
   assert (CastWidth <= IVWidth && "Unexpected width while widening compare.");
 
   // Widen the compare instruction.
   IRBuilder<> Builder(
       getInsertPointForUses(DU.NarrowUse, DU.NarrowDef, DT, LI));
   DU.NarrowUse->replaceUsesOfWith(DU.NarrowDef, DU.WideDef);
 
   // Widen the other operand of the compare, if necessary.
   if (CastWidth < IVWidth) {
     Value *ExtOp = createExtendInst(Op, WideType, Cmp->isSigned(), Cmp);
     DU.NarrowUse->replaceUsesOfWith(Op, ExtOp);
   }
   return true;
 }
 
 /// Determine whether an individual user of the narrow IV can be widened. If so,
 /// return the wide clone of the user.
 Instruction *WidenIV::widenIVUse(NarrowIVDefUse DU, SCEVExpander &Rewriter) {
+  assert(ExtendKindMap.count(DU.NarrowDef) &&
+         "Should already know the kind of extension used to widen NarrowDef");
 
   // Stop traversing the def-use chain at inner-loop phis or post-loop phis.
   if (PHINode *UsePhi = dyn_cast<PHINode>(DU.NarrowUse)) {
     if (LI->getLoopFor(UsePhi->getParent()) != L) {
       // For LCSSA phis, sink the truncate outside the loop.
       // After SimplifyCFG most loop exit targets have a single predecessor.
       // Otherwise fall back to a truncate within the loop.
       if (UsePhi->getNumOperands() != 1)
if (LI->getLoopFor(UsePhi->getParent()) != L) {
         UsePhi->replaceAllUsesWith(Trunc);
         DeadInsts.emplace_back(UsePhi);
         DEBUG(dbgs() << "INDVARS: Widen lcssa phi " << *UsePhi
               << " to " << *WidePhi << "\n");
       }
       return nullptr;
     }
   }
+
+  // This narrow use can be widened by a sext if it's non-negative or its narrow
+  // def was widended by a sext. Same for zext.
+  auto canWidenBySExt = [&]() {
+    return DU.NeverNegative || ExtendKindMap[DU.NarrowDef] == SignExtended;
+  };
+  auto canWidenByZExt = [&]() {
+    return DU.NeverNegative || ExtendKindMap[DU.NarrowDef] == ZeroExtended;
+  };
+
   // Our raison d'etre! Eliminate sign and zero extension.
-  if (IsSigned ? isa<SExtInst>(DU.NarrowUse) : isa<ZExtInst>(DU.NarrowUse)) {
+  if ((isa<SExtInst>(DU.NarrowUse) && canWidenBySExt()) ||
+      (isa<ZExtInst>(DU.NarrowUse) && canWidenByZExt())) {
     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 ((isa<SExtInst>(DU.NarrowUse) && canWidenBySExt()) ||
     // new loop phi. If we preserved IVUsers analysis, we would also want to
     // push the uses of WideDef here.
 
     // 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)
+  WidenedRecTy WideAddRec = getWideRecurrence(DU);
+  if (!WideAddRec.first)
     WideAddRec = getExtendedOperandRecurrence(DU);
 
-  if (!WideAddRec) {
+  assert((WideAddRec.first == nullptr) == (WideAddRec.second == Unknown));
+  if (!WideAddRec.first) {
     // 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,
     // eventually isolating the original narrow IV so it can be removed.
     truncateIVUse(DU, DT, LI);
     return nullptr;
   }
   // Assume block terminators cannot evaluate to a recurrence. We can't to
   // insert a Trunc after a terminator if there happens to be a critical edge.
   assert(DU.NarrowUse != DU.NarrowUse->getParent()->getTerminator() &&
          "SCEV is not expected to evaluate a block terminator");
 
   // Reuse the IV increment that SCEVExpander created as long as it dominates
   // NarrowUse.
   Instruction *WideUse = nullptr;
-  if (WideAddRec == WideIncExpr && Rewriter.hoistIVInc(WideInc, DU.NarrowUse))
+  if (WideAddRec.first == WideIncExpr && 
+      Rewriter.hoistIVInc(WideInc, DU.NarrowUse))
     WideUse = WideInc;
   else {
-    WideUse = cloneIVUser(DU, WideAddRec);
+    WideUse = cloneIVUser(DU, WideAddRec.first);
     if (!WideUse)
       return nullptr;
   }
   // Evaluation of WideAddRec ensured that the narrow expression could be
   // extended outside the loop without overflow. This suggests that the wide use
   // evaluates to the same expression as the extended narrow use, but doesn't
   // absolutely guarantee it. Hence the following failsafe check. In rare cases
   // where it fails, we simply throw away the newly created wide use.
-  if (WideAddRec != SE->getSCEV(WideUse)) {
+  if (WideAddRec.first != SE->getSCEV(WideUse)) {
     DEBUG(dbgs() << "Wide use expression mismatch: " << *WideUse
-          << ": " << *SE->getSCEV(WideUse) << " != " << *WideAddRec << "\n");
+          << ": " << *SE->getSCEV(WideUse) << " != " << *WideAddRec.first << "\n");
     DeadInsts.emplace_back(WideUse);
     return nullptr;
   }
 
+  ExtendKindMap[DU.NarrowUse] = WideAddRec.second;
   // 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) {
   // Is this phi an induction variable?
   const SCEVAddRecExpr *AddRec = dyn_cast<SCEVAddRecExpr>(SE->getSCEV(OrigPhi));
   if (!AddRec)
     return nullptr;
 
   // Widen the induction variable expression.
-  const SCEV *WideIVExpr = IsSigned ?
-    SE->getSignExtendExpr(AddRec, WideType) :
-    SE->getZeroExtendExpr(AddRec, WideType);
+  const SCEV *WideIVExpr = ExtendKindMap[OrigPhi] == SignExtended
+                               ? SE->getSignExtendExpr(AddRec, WideType)
+                               : SE->getZeroExtendExpr(AddRec, WideType);
 
   assert(SE->getEffectiveSCEVType(WideIVExpr->getType()) == WideType &&
          "Expect the new IV expression to preserve its type");
 
   // Can the IV be extended outside the loop without overflow?
   AddRec = dyn_cast<SCEVAddRecExpr>(WideIVExpr);
   if (!AddRec || AddRec->getLoop() != L)
     return nullptr;
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.

test/Transforms/IndVarSimplify/iv-widen-elim-ext.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -indvars -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-p:64:64:64-n8:16:32:64-S128"
+
+; When widening IV and its users, trunc and zext/sext are not needed
+; if the original 32-bit user is known to be non-negative, whether
+; the IV is considered signed or unsigned.
+define void @foo(i32* %A, i32* %B, i32* %C, i32 %N) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 0, %N
+; CHECK-NEXT:    br i1 [[CMP1]], label %for.body.lr.ph, label %for.end
+; CHECK:       for.body.lr.ph:
+; CHECK-NEXT:    br label %for.body
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV:%.*]].next, %for.inc ], [ 0, %for.body.lr.ph ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* %B, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 2
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds i32, i32* %C, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[TMP0]], [[TMP2]]
+; CHECK-NEXT:    [[ARRAYIDX5:%.*]] = getelementptr inbounds i32, i32* %A, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    store i32 [[ADD3]], i32* [[ARRAYIDX5]], align 4
+; CHECK-NEXT:    br label %for.inc
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext i32 %N to i64
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT]], [[WIDE_TRIP_COUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label %for.body, label %for.cond.for.end_crit_edge
+; CHECK:       for.cond.for.end_crit_edge:
+; CHECK-NEXT:    br label %for.end
+; CHECK:       for.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp1 = icmp slt i32 0, %N
+  br i1 %cmp1, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.inc
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  %idxprom = sext i32 %i.02 to i64
+  %arrayidx = getelementptr inbounds i32, i32* %B, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  %add = add nsw i32 %i.02, 2
+  %idxprom1 = zext i32 %add to i64
+  %arrayidx2 = getelementptr inbounds i32, i32* %C, i64 %idxprom1
+  %1 = load i32, i32* %arrayidx2, align 4
+  %add3 = add nsw i32 %0, %1
+  %idxprom4 = zext i32 %i.02 to i64
+  %arrayidx5 = getelementptr inbounds i32, i32* %A, i64 %idxprom4
+  store i32 %add3, i32* %arrayidx5, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %N
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.inc
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  ret void
+}
+
+define void @foo1(i32* %A, i32* %B, i32* %C, i32 %N) {
+; CHECK-LABEL: @foo1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 0, %N
+; CHECK-NEXT:    br i1 [[CMP1]], label %for.body.lr.ph, label %for.end
+; CHECK:       for.body.lr.ph:
+; CHECK-NEXT:    br label %for.body
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV:%.*]].next, %for.inc ], [ 0, %for.body.lr.ph ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds i32, i32* %B, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP0:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[TMP1:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 2
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds i32, i32* %C, i64 [[TMP1]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[ADD3:%.*]] = add nsw i32 [[TMP0]], [[TMP2]]
+; CHECK-NEXT:    [[ARRAYIDX5:%.*]] = getelementptr inbounds i32, i32* %A, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    store i32 [[ADD3]], i32* [[ARRAYIDX5]], align 4
+; CHECK-NEXT:    br label %for.inc
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext i32 %N to i64
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT]], [[WIDE_TRIP_COUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label %for.body, label %for.cond.for.end_crit_edge
+; CHECK:       for.cond.for.end_crit_edge:
+; CHECK-NEXT:    br label %for.end
+; CHECK:       for.end:
+; CHECK-NEXT:    ret void
+;
+entry:
+  %cmp1 = icmp slt i32 0, %N
+  br i1 %cmp1, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.inc
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  %idxprom = zext i32 %i.02 to i64
+  %arrayidx = getelementptr inbounds i32, i32* %B, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  %add = add nsw i32 %i.02, 2
+  %idxprom1 = sext i32 %add to i64
+  %arrayidx2 = getelementptr inbounds i32, i32* %C, i64 %idxprom1
+  %1 = load i32, i32* %arrayidx2, align 4
+  %add3 = add nsw i32 %0, %1
+  %idxprom4 = sext i32 %i.02 to i64
+  %arrayidx5 = getelementptr inbounds i32, i32* %A, i64 %idxprom4
+  store i32 %add3, i32* %arrayidx5, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %N
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.inc
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  ret void
+}
+
+
+@a = common global [100 x i32] zeroinitializer, align 16
+@b = common global [100 x i32] zeroinitializer, align 16
+
+define i32 @foo2(i32 %M) {
+; CHECK-LABEL: @foo2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 0, %M
+; CHECK-NEXT:    br i1 [[CMP1]], label %for.body.lr.ph, label %for.end
+; CHECK:       for.body.lr.ph:
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 %M to i64
+; CHECK-NEXT:    br label %for.body
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV:%.*]].next, %for.inc ], [ 0, %for.body.lr.ph ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @b, i64 0, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP3:%.*]] = add nsw i64 [[INDVARS_IV]], [[TMP0]]
+; CHECK-NEXT:    [[ARRAYIDX5:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 [[TMP3]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[ARRAYIDX5]], align 4
+; CHECK-NEXT:    br label %for.inc
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext i32 %M to i64
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT]], [[WIDE_TRIP_COUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label %for.body, label %for.cond.for.end_crit_edge
+; CHECK:       for.cond.for.end_crit_edge:
+; CHECK-NEXT:    br label %for.end
+; CHECK:       for.end:
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @dummy(i32* getelementptr inbounds ([100 x i32], [100 x i32]* @a, i32 0, i32 0), i32* getelementptr inbounds ([100 x i32], [100 x i32]* @b, i32 0, i32 0))
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %cmp1 = icmp slt i32 0, %M
+  br i1 %cmp1, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.inc
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  %idxprom = zext i32 %i.02 to i64
+  %arrayidx = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  %idxprom1 = sext i32 %i.02 to i64
+  %arrayidx2 = getelementptr inbounds [100 x i32], [100 x i32]* @b, i64 0, i64 %idxprom1
+  %1 = load i32, i32* %arrayidx2, align 4
+  %add = add nsw i32 %0, %1
+  %add3 = add nsw i32 %i.02, %M
+  %idxprom4 = sext i32 %add3 to i64
+  %arrayidx5 = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 %idxprom4
+  store i32 %add, i32* %arrayidx5, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %M
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.inc
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  %call = call i32 @dummy(i32* getelementptr inbounds ([100 x i32], [100 x i32]* @a, i32 0, i32 0), i32* getelementptr inbounds ([100 x i32], [100 x i32]* @b, i32 0, i32 0))
+  ret i32 0
+}
+
+declare i32 @dummy(i32*, i32*)
+
+; A case where zext should not be eliminated when its operands could only be extended by sext.
+define i32 @foo3(i32 %M) {
+; CHECK-LABEL: @foo3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp slt i32 0, %M
+; CHECK-NEXT:    br i1 [[CMP1]], label %for.body.lr.ph, label %for.end
+; CHECK:       for.body.lr.ph:
+; CHECK-NEXT:    [[TMP0:%.*]] = sext i32 %M to i64
+; CHECK-NEXT:    br label %for.body
+; CHECK:       for.body:
+; CHECK-NEXT:    [[INDVARS_IV:%.*]] = phi i64 [ [[INDVARS_IV:%.*]].next, %for.inc ], [ 0, %for.body.lr.ph ]
+; CHECK-NEXT:    [[ARRAYIDX:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, i32* [[ARRAYIDX]], align 4
+; CHECK-NEXT:    [[ARRAYIDX2:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @b, i64 0, i64 [[INDVARS_IV]]
+; CHECK-NEXT:    [[TMP2:%.*]] = load i32, i32* [[ARRAYIDX2]], align 4
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    [[TMP3:%.*]] = add nsw i64 [[INDVARS_IV]], [[TMP0]]
+; CHECK-NEXT:    [[TMP4:%.*]] = trunc i64 [[TMP3]] to i32
+; CHECK-NEXT:    [[IDXPROM4:%.*]] = zext i32 [[TMP4]] to i64
+; CHECK-NEXT:    [[ARRAYIDX5:%.*]] = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 [[IDXPROM4]]
+; CHECK-NEXT:    store i32 [[ADD]], i32* [[ARRAYIDX5]], align 4
+; CHECK-NEXT:    br label %for.inc
+; CHECK:       for.inc:
+; CHECK-NEXT:    [[INDVARS_IV_NEXT:%.*]] = add nuw nsw i64 [[INDVARS_IV]], 1
+; CHECK-NEXT:    [[WIDE_TRIP_COUNT:%.*]] = zext i32 %M to i64
+; CHECK-NEXT:    [[EXITCOND:%.*]] = icmp ne i64 [[INDVARS_IV_NEXT]], [[WIDE_TRIP_COUNT]]
+; CHECK-NEXT:    br i1 [[EXITCOND]], label %for.body, label %for.cond.for.end_crit_edge
+; CHECK:       for.cond.for.end_crit_edge:
+; CHECK-NEXT:    br label %for.end
+; CHECK:       for.end:
+; CHECK-NEXT:    [[CALL:%.*]] = call i32 @dummy(i32* getelementptr inbounds ([100 x i32], [100 x i32]* @a, i32 0, i32 0), i32* getelementptr inbounds ([100 x i32], [100 x i32]* @b, i32 0, i32 0))
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %cmp1 = icmp slt i32 0, %M
+  br i1 %cmp1, label %for.body.lr.ph, label %for.end
+
+for.body.lr.ph:                                   ; preds = %entry
+  br label %for.body
+
+for.body:                                         ; preds = %for.body.lr.ph, %for.inc
+  %i.02 = phi i32 [ 0, %for.body.lr.ph ], [ %inc, %for.inc ]
+  %idxprom = sext i32 %i.02 to i64
+  %arrayidx = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 %idxprom
+  %0 = load i32, i32* %arrayidx, align 4
+  %idxprom1 = sext i32 %i.02 to i64
+  %arrayidx2 = getelementptr inbounds [100 x i32], [100 x i32]* @b, i64 0, i64 %idxprom1
+  %1 = load i32, i32* %arrayidx2, align 4
+  %add = add nsw i32 %0, %1
+  %add3 = add nsw i32 %i.02, %M
+  %idxprom4 = zext i32 %add3 to i64
+  %arrayidx5 = getelementptr inbounds [100 x i32], [100 x i32]* @a, i64 0, i64 %idxprom4
+  store i32 %add, i32* %arrayidx5, align 4
+  br label %for.inc
+
+for.inc:                                          ; preds = %for.body
+  %inc = add nsw i32 %i.02, 1
+  %cmp = icmp slt i32 %inc, %M
+  br i1 %cmp, label %for.body, label %for.cond.for.end_crit_edge
+
+for.cond.for.end_crit_edge:                       ; preds = %for.inc
+  br label %for.end
+
+for.end:                                          ; preds = %for.cond.for.end_crit_edge, %entry
+  %call = call i32 @dummy(i32* getelementptr inbounds ([100 x i32], [100 x i32]* @a, i32 0, i32 0), i32* getelementptr inbounds ([100 x i32], [100 x i32]* @b, i32 0, i32 0))
+  ret i32 0
+}