[LSR] Generate and use zero extends

[sanjoy]

Summary

If a scale or a base register can be rewritten as "Zext({A,+,1})" then
LSR will now consider a formula of that form in its normal cost
computation.

Depends on D9180

[sanjoy]

ping!

[atrick]

This looks reasonable. Thanks!

Does your test fail on trunk? This is what the test outputs for me:

ok_161: ; preds = %ok_158

%lsr.iv.next = add nuw nsw i64 %lsr.iv, 1
%4 = add i64 %0, %lsr.iv.next
%tmp1 = trunc i64 %4 to i32
%tmp188 = icmp slt i32 %tmp1, %tmp160
br i1 %tmp188, label %ok_146, label %block_81

Would you be able to test both post-inc and pre-inc variants? We've a lot of bugs because of TransformForPostIncUse.

lib/Transforms/Scalar/LoopStrengthReduce.cpp

3685 ↗

(On Diff #24191)

Please comment this at least as well as your commit comment if not better.

[sanjoy]

Does your test fail on trunk? This is what the test outputs for me:

ok_161: ; preds = %ok_158

%lsr.iv.next = add nuw nsw i64 %lsr.iv, 1
%4 = add i64 %0, %lsr.iv.next
%tmp1 = trunc i64 %4 to i32
%tmp188 = icmp slt i32 %tmp1, %tmp160
br i1 %tmp188, label %ok_146, label %block_81

That's behavior I'm trying to avoid -- I don't want two additions on
every iteration of the loop.

To elaborate a little bit, here is the -debug-only=loop-reduce on the
test case before the change:

'''
After filtering out undesirable candidates:
LSR is examining the following uses:

LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i32
  reg({%tmp156,+,1}<nuw><nsw><%ok_146>)
  reg((zext i32 %tmp156 to i64)) + 1*reg({2,+,1}<nw><%ok_146>) + imm(-2)
  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>)
  reg((zext i32 %tmp156 to i64)) + 1*reg({-1,+,1}<nw><%ok_146>) + imm(1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({3,+,1}<nw><%ok_146>) + imm(-3)
LSR Use: Kind=Basic, Offsets={0}, all-fixups-outside-loop, widest

fixup type: i64

  reg({(-1 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  reg({(zext i32 %tmp156 to i64),+,1}<nuw><nsw><%ok_146>) + imm(-1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>) + imm(-1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({-1,+,1}<nw><%ok_146>)
  reg((-1 + (zext i32 %tmp156 to i64))) + 1*reg({0,+,1}<nuw><%ok_146>)
  reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>) + imm(-4)
  reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>) + imm(-3)
LSR Use: Kind=Address of double, Offsets={16}, widest fixup type: double*
  -16 + reg({(16 + (8 * (zext i32 %tmp156 to i64)) + %d),+,8}<nw><%ok_146>)
  -24 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 4*reg({6,+,2}<%ok_146>)
  -16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 2*reg({8,+,4}<%ok_146>)
  -24 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 2*reg({12,+,4}<%ok_146>)
  -16 + reg(%d) + 8*reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  -24 + reg(%d) + 8*reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  -8 + reg((16 + %d)<nsw>) + 8*reg({(-1 + (zext i32 %tmp156 to

i64)),+,1}<nw><%ok_146>)

-16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 4*reg({4,+,2}<%ok_146>)
-16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) +

8*reg({2,+,1}<nw><%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

2*reg({0,+,4}<%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

4*reg({0,+,2}<%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

8*reg({0,+,1}<nuw><%ok_146>)

-16 + reg((16 + %d)<nsw>) + 2*reg({(4 * (zext i32 %tmp156 to

i64)),+,4}<nuw><nsw><%ok_146>)

-16 + reg((16 + %d)<nsw>) + 4*reg({(2 * (zext i32 %tmp156 to

i64)),+,2}<%ok_146>)

-16 + reg((16 + %d)<nsw>) + 8*reg({(zext i32 %tmp156 to

i64),+,1}<nuw><nsw><%ok_146>)

LSR Use: Kind=Address of i32, Offsets={12}, widest fixup type: i32*
  -12 + reg({(12 + (4 * (zext i32 %tmp156 to i64)) + %b),+,4}<nw><%ok_146>)
  -20 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

2*reg({4,+,2}<%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

1*reg({0,+,4}<%ok_146>)

-12 + reg((12 + %b)<nsw>) + 1*reg({(4 * (zext i32 %tmp156 to

i64)),+,4}<nuw><nsw><%ok_146>)

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) + 2*reg({6,+,2}<%ok_146>)
-20 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

1*reg({8,+,4}<%ok_146>)

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) + 1*reg({12,+,4}<%ok_146>)
-8 + reg(%b) + 4*reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
-12 + reg(%b) + 4*reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
-8 + reg((12 + %b)<nsw>) + 4*reg({(-1 + (zext i32 %tmp156 to

i64)),+,1}<nw><%ok_146>)

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) +

4*reg({3,+,1}<nw><%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

2*reg({0,+,2}<%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

4*reg({0,+,1}<nuw><%ok_146>)

-12 + reg((12 + %b)<nsw>) + 2*reg({(2 * (zext i32 %tmp156 to

i64)),+,2}<%ok_146>)

-12 + reg((12 + %b)<nsw>) + 4*reg({(zext i32 %tmp156 to

i64),+,1}<nuw><nsw><%ok_146>)
New best at 4 regs, with addrec cost 2, plus 2 scale cost, plus 9 imm
cost, plus 3 setup cost.
Regs: {%tmp156,+,1}<nuw><nsw><%ok_146> {(-1 + (zext i32 %tmp156 to
i64)),+,1}<nw><%ok_146> (16 + %d)<nsw> (12 + %b)<nsw>
New best at 4 regs, with addrec cost 1, plus 3 base adds, plus 2 scale
cost, plus 3 setup cost.
Regs: {0,+,1}<nuw><%ok_146> (zext i32 %tmp156 to i64) (16 + (8 *
(zext i32 %tmp156 to i64)) + %d) (12 + (4 * (zext i32 %tmp156 to i64))
+ %b)

The chosen solution requires 4 regs, with addrec cost 1, plus 3 base
adds, plus 2 scale cost, plus 3 setup cost:

LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i32
  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>)
LSR Use: Kind=Basic, Offsets={0}, all-fixups-outside-loop, widest

fixup type: i64

  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>) + imm(-1)
LSR Use: Kind=Address of double, Offsets={16}, widest fixup type: double*
  -16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

8*reg({0,+,1}<nuw><%ok_146>)

LSR Use: Kind=Address of i32, Offsets={12}, widest fixup type: i32*
  -12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

4*reg({0,+,1}<nuw><%ok_146>)
'''

And here is the debug output after the change. LSR now "sees" a few
more candidate formulae (marked with '<- new solution +++++'):

'''
After filtering out undesirable candidates:
LSR is examining the following uses:

LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i32
  reg({%tmp156,+,1}<nuw><nsw><%ok_146>)
  reg(%tmp156) + 1*reg({0,+,1}<nuw><%ok_146>)
  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>)
  reg((zext i32 %tmp156 to i64)) + 1*reg({-1,+,1}<nw><%ok_146>) + imm(1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({3,+,1}<nw><%ok_146>) + imm(-3)
  reg((zext i32 %tmp156 to i64)) + 1*reg({2,+,1}<nw><%ok_146>) + imm(-2)
LSR Use: Kind=Basic, Offsets={0}, all-fixups-outside-loop, widest

fixup type: i64

  reg({(-1 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  reg({(zext i32 %tmp156 to i64),+,1}<nuw><nsw><%ok_146>) + imm(-1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({0,+,1}<nuw><%ok_146>) + imm(-1)
  reg((zext i32 %tmp156 to i64)) + 1*reg({-1,+,1}<nw><%ok_146>)
  reg((-1 + (zext i32 %tmp156 to i64))) + 1*reg({0,+,1}<nuw><%ok_146>)
  reg({%tmp156,+,1}<nuw><nsw><%ok_146>) + imm(-1)   <- new solution +++++
  reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>) + imm(-4)
  reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>) + imm(-3)
LSR Use: Kind=Address of double, Offsets={16}, widest fixup type: double*
  -16 + reg({(16 + (8 * (zext i32 %tmp156 to i64)) + %d),+,8}<nw><%ok_146>)
  -24 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 4*reg({6,+,2}<%ok_146>)
  -16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 2*reg({8,+,4}<%ok_146>)
  -24 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 2*reg({12,+,4}<%ok_146>)
  -16 + reg(%d) + 8*reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  -24 + reg(%d) + 8*reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
  -8 + reg((16 + %d)<nsw>) + 8*reg({(-1 + (zext i32 %tmp156 to

i64)),+,1}<nw><%ok_146>)

-16 + reg((16 + %d)<nsw>) +

8*reg({%tmp156,+,1}<nuw><nsw><%ok_146>) <- new solution +++++

-16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) + 4*reg({4,+,2}<%ok_146>)
-16 + reg(((8 * (zext i32 %tmp156 to i64)) + %d)) +

8*reg({2,+,1}<nw><%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

2*reg({0,+,4}<%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

4*reg({0,+,2}<%ok_146>)

-16 + reg((16 + (8 * (zext i32 %tmp156 to i64)) + %d)) +

8*reg({0,+,1}<nuw><%ok_146>)

-16 + reg((16 + %d)<nsw>) + 2*reg({(4 * (zext i32 %tmp156 to

i64)),+,4}<nuw><nsw><%ok_146>)

-16 + reg((16 + %d)<nsw>) + 4*reg({(2 * (zext i32 %tmp156 to

i64)),+,2}<%ok_146>)

-16 + reg((16 + %d)<nsw>) + 8*reg({(zext i32 %tmp156 to

i64),+,1}<nuw><nsw><%ok_146>)

LSR Use: Kind=Address of i32, Offsets={12}, widest fixup type: i32*
  -12 + reg({(12 + (4 * (zext i32 %tmp156 to i64)) + %b),+,4}<nw><%ok_146>)
  -20 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

2*reg({4,+,2}<%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

1*reg({0,+,4}<%ok_146>)

-12 + reg((12 + %b)<nsw>) + 1*reg({(4 * (zext i32 %tmp156 to

i64)),+,4}<nuw><nsw><%ok_146>)

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) + 2*reg({6,+,2}<%ok_146>)
-20 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

1*reg({8,+,4}<%ok_146>)

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) + 1*reg({12,+,4}<%ok_146>)
-8 + reg(%b) + 4*reg({(2 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
-12 + reg(%b) + 4*reg({(3 + (zext i32 %tmp156 to i64)),+,1}<nw><%ok_146>)
-8 + reg((12 + %b)<nsw>) + 4*reg({(-1 + (zext i32 %tmp156 to

i64)),+,1}<nw><%ok_146>)

-12 + reg((12 + %b)<nsw>) +

4*reg({%tmp156,+,1}<nuw><nsw><%ok_146>) <- new solution +++++

-12 + reg(((4 * (zext i32 %tmp156 to i64)) + %b)) +

4*reg({3,+,1}<nw><%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

2*reg({0,+,2}<%ok_146>)

-12 + reg((12 + (4 * (zext i32 %tmp156 to i64)) + %b)) +

4*reg({0,+,1}<nuw><%ok_146>)

-12 + reg((12 + %b)<nsw>) + 2*reg({(2 * (zext i32 %tmp156 to

i64)),+,2}<%ok_146>)

-12 + reg((12 + %b)<nsw>) + 4*reg({(zext i32 %tmp156 to

i64),+,1}<nuw><nsw><%ok_146>)
New best at 3 regs, with addrec cost 1, plus 1 base add, plus 2 scale
cost, plus 2 setup cost.
Regs: {%tmp156,+,1}<nuw><nsw><%ok_146> (16 + %d)<nsw> (12 + %b)<nsw>

The chosen solution requires 3 regs, with addrec cost 1, plus 1 base
add, plus 2 scale cost, plus 2 setup cost:

LSR Use: Kind=Basic, Offsets={0}, widest fixup type: i32
  reg({%tmp156,+,1}<nuw><nsw><%ok_146>)
LSR Use: Kind=Basic, Offsets={0}, all-fixups-outside-loop, widest

fixup type: i64

  reg({%tmp156,+,1}<nuw><nsw><%ok_146>) + imm(-1)
LSR Use: Kind=Address of double, Offsets={16}, widest fixup type: double*
  -16 + reg((16 + %d)<nsw>) + 8*reg({%tmp156,+,1}<nuw><nsw><%ok_146>)
LSR Use: Kind=Address of i32, Offsets={12}, widest fixup type: i32*
  -12 + reg((12 + %b)<nsw>) + 4*reg({%tmp156,+,1}<nuw><nsw><%ok_146>)

'''

The new formulae lets LSR choose a better solution.

Looking at the output tells me that I should definitely add something
to the textual representation of a formulae that has an implicit zero
extend.

Would you be able to test both post-inc and pre-inc variants? We've a
lot of bugs because of TransformForPostIncUse.

I think that is a good idea. I'm not familiar with LSR so I suspect
it will take me some time to get there, but just that is actually a
good reason for writing these tests. :)

It also looks like (I'm not a 100% sure) that there's an inconsistency
in how LSR treats overflow in post-inc expressions vs. how SCEV treats
them. For instance, in ScalarEvolution.cpp line ~1500 we say that
{S,+,X} is <nuw> if "S + MaxBECount * X" does not unsigned
overflow. But that does not say anything about the overflow of the
post-inc SCEV expression (e.g. if BECount == 0 then all add recs are
<nuw>).

The problem is that TransformForPostIncUse(zext({S,+,X})) =
zext({S+X,+,X}) -- ScalarEvolutionNormalization.cpp, line ~100.

Now if {S,+,1} is nuw then {zext S,+,1} = zext({S,+,1}), but
{(zext S) + 1,+,1} is != zext({S+1,+,1}) by SCEV's definition of nuw.
But TransformForPostIncUse({zext S,+,1}) =
TransformForPostIncUse(zext({S,+,1})) = zext({S+1,+,1}) when it should
really be {(zext S) + 1,+,1}. I don't have a concrete example where
this is a problem, but this definitely looks fishy to me.

• Sanjoy

[sanjoy]

address review:

• change Formula::print to indicate that a specific formula is a zext formula
• add comment on LSRInstance::GenerateZExts
• the existing test case tests both pre-inc and post-inc uses (commented). Please let me know if this is not what you meant.

[sanjoy]

ping!

[qcolombet]

FWIW, LGTM.

I let Andy comment on the latest changes.

Thanks,
Q.

[sanjoy]

ping!

[sanjoy]

ping!

[sanjoy]

ping!

[atrick]

LGTM. Sorry to neglect this one that was clearly blocked on me.

llvm/trunk/lib/Transforms/Scalar/LoopStrengthReduce.cpp

struct Formula {
   /// when Scale is not zero.
   const SCEV *ScaledReg;
 
   /// UnfoldedOffset - An additional constant offset which added near the
   /// use. This requires a temporary register, but the offset itself can
   /// live in an add immediate field rather than a register.
   int64_t UnfoldedOffset;
 
+  /// ZeroExtendScaledReg - This formula zero extends the scale register to
+  /// ZeroExtendType before its use.
+  bool ZeroExtendScaledReg;
+
+  /// ZeroExtendBaseReg - This formula zero extends all the base registers to
+  /// ZeroExtendType before their use.
+  bool ZeroExtendBaseReg;
+
+  /// ZeroExtendType - The destination type of the zero extension implied by
+  /// the above two booleans.
+  Type *ZeroExtendType;
+
   Formula()
       : BaseGV(nullptr), BaseOffset(0), HasBaseReg(false), Scale(0),
-        ScaledReg(nullptr), UnfoldedOffset(0) {}
+        ScaledReg(nullptr), UnfoldedOffset(0), ZeroExtendScaledReg(false),
+        ZeroExtendBaseReg(false), ZeroExtendType(nullptr) {}
 
   void InitialMatch(const SCEV *S, Loop *L, ScalarEvolution &SE);
 
   bool isCanonical() const;
 
   void Canonicalize();
 
   bool Unscale();

 /// addrec strides.
 size_t Formula::getNumRegs() const {
   return !!ScaledReg + BaseRegs.size();
 }
 
 /// getType - Return the type of this formula, if it has one, or null
 /// otherwise. This type is meaningless except for the bit size.
 Type *Formula::getType() const {
-  return !BaseRegs.empty() ? BaseRegs.front()->getType() :
-         ScaledReg ? ScaledReg->getType() :
-         BaseGV ? BaseGV->getType() :
-         nullptr;
+  return ZeroExtendType
+             ? ZeroExtendType
+             : !BaseRegs.empty()
+                   ? BaseRegs.front()->getType()
+                   : ScaledReg ? ScaledReg->getType()
+                               : BaseGV ? BaseGV->getType() : nullptr;
 }
 
 /// DeleteBaseReg - Delete the given base reg from the BaseRegs list.
 void Formula::DeleteBaseReg(const SCEV *&S) {
   if (&S != &BaseRegs.back())
     std::swap(S, BaseRegs.back());
   BaseRegs.pop_back();
 }
if (BaseGV) {
     BaseGV->printAsOperand(OS, /*PrintType=*/false);
   }
   if (BaseOffset != 0) {
     if (!First) OS << " + "; else First = false;
     OS << BaseOffset;
   }
   for (const SCEV *BaseReg : BaseRegs) {
     if (!First) OS << " + "; else First = false;
+    if (ZeroExtendBaseReg)
+      OS << "reg(zext " << *BaseReg << " to " << *ZeroExtendType << ')';
+    else
       OS << "reg(" << *BaseReg << ')';
   }
   if (HasBaseReg && BaseRegs.empty()) {
     if (!First) OS << " + "; else First = false;
     OS << "**error: HasBaseReg**";
   } else if (!HasBaseReg && !BaseRegs.empty()) {
     if (!First) OS << " + "; else First = false;
     OS << "**error: !HasBaseReg**";
   }
   if (Scale != 0) {
     if (!First) OS << " + "; else First = false;
     OS << Scale << "*reg(";
-    if (ScaledReg)
-      OS << *ScaledReg;
+    if (ScaledReg) {
+      if (ZeroExtendScaledReg)
+        OS << "(zext " << *ScaledReg << " to " << *ZeroExtendType << ')';
       else
+        OS << *ScaledReg;
+    } else
       OS << "<unknown>";
     OS << ')';
   }
   if (UnfoldedOffset != 0) {
     if (!First) OS << " + ";
     OS << "imm(" << UnfoldedOffset << ')';
   }
 }
class LSRInstance {
   void GenerateConstantOffsetsImpl(LSRUse &LU, unsigned LUIdx,
                                    const Formula &Base,
                                    const SmallVectorImpl<int64_t> &Worklist,
                                    size_t Idx, bool IsScaledReg = false);
   void GenerateConstantOffsets(LSRUse &LU, unsigned LUIdx, Formula Base);
   void GenerateICmpZeroScales(LSRUse &LU, unsigned LUIdx, Formula Base);
   void GenerateScales(LSRUse &LU, unsigned LUIdx, Formula Base);
   void GenerateTruncates(LSRUse &LU, unsigned LUIdx, Formula Base);
+  void GenerateZExts(LSRUse &LU, unsigned LUIdx, Formula Base);
   void GenerateCrossUseConstantOffsets();
   void GenerateAllReuseFormulae();
 
   void FilterOutUndesirableDedicatedRegisters();
 
   size_t EstimateSearchSpaceComplexity() const;
   void NarrowSearchSpaceByDetectingSupersets();
   void NarrowSearchSpaceByCollapsingUnrolledCode();
if (SrcTy != DstTy && TTI.isTruncateFree(SrcTy, DstTy)) {
       if (!F.hasRegsUsedByUsesOtherThan(LUIdx, RegUses))
         continue;
 
       (void)InsertFormula(LU, LUIdx, F);
     }
   }
 }
 
+/// GenerateZExts - If a scale or a base register can be rewritten as
+/// "Zext({A,+,1})" then consider a formula of that form.
+void LSRInstance::GenerateZExts(LSRUse &LU, unsigned LUIdx, Formula Base) {
+  // Don't bother with symbolic values.
+  if (Base.BaseGV)
+    return;
+
+  auto CanBeNarrowed = [&](const SCEV *Reg) -> const SCEV * {
+    // Check if the register is an increment can be rewritten as zext(R) where
+    // the zext is free.
+
+    const auto *RegAR = dyn_cast_or_null<SCEVAddRecExpr>(Reg);
+    if (!RegAR)
+      return nullptr;
+
+    const auto *ZExtStart = dyn_cast<SCEVZeroExtendExpr>(RegAR->getStart());
+    const auto *ConstStep =
+        dyn_cast<SCEVConstant>(RegAR->getStepRecurrence(SE));
+    if (!ZExtStart || !ConstStep || ConstStep->getValue()->getValue() != 1)
+      return nullptr;
+
+    const SCEV *NarrowStart = ZExtStart->getOperand();
+    if (!TTI.isZExtFree(NarrowStart->getType(), ZExtStart->getType()))
+      return nullptr;
+
+    const auto *NarrowAR = dyn_cast<SCEVAddRecExpr>(
+        SE.getAddRecExpr(NarrowStart, SE.getConstant(NarrowStart->getType(), 1),
+                         RegAR->getLoop(), RegAR->getNoWrapFlags()));
+
+    if (!NarrowAR || !NarrowAR->getNoWrapFlags(SCEV::FlagNUW))
+      return nullptr;
+
+    return NarrowAR;
+  };
+
+  if (Base.ScaledReg && !Base.ZeroExtendType)
+    if (const SCEV *S = CanBeNarrowed(Base.ScaledReg)) {
+      Formula F = Base;
+      F.ZeroExtendType = Base.ScaledReg->getType();
+      F.ZeroExtendScaledReg = true;
+      F.ScaledReg = S;
+
+      if (isLegalUse(TTI, LU.MinOffset, LU.MaxOffset, LU.Kind, LU.AccessTy, F))
+        InsertFormula(LU, LUIdx, F);
+    }
+
+  if (Base.BaseRegs.size() == 1 && !Base.ZeroExtendType)
+    if (const SCEV *S = CanBeNarrowed(Base.BaseRegs[0])) {
+      Formula F = Base;
+      F.ZeroExtendType = Base.BaseRegs[0]->getType();
+      F.ZeroExtendBaseReg = true;
+      F.BaseRegs[0] = S;
+
+      if (isLegalUse(TTI, LU.MinOffset, LU.MaxOffset, LU.Kind, LU.AccessTy, F))
+        InsertFormula(LU, LUIdx, F);
+    }
+}
+
 namespace {
 
 /// WorkItem - Helper class for GenerateCrossUseConstantOffsets. It's used to
 /// defer modifications so that the search phase doesn't have to worry about
 /// the data structures moving underneath it.
 struct WorkItem {
   size_t LUIdx;
   int64_t Imm;
for (size_t i = 0, f = LU.Formulae.size(); i != f; ++i)
       GenerateICmpZeroScales(LU, LUIdx, LU.Formulae[i]);
     for (size_t i = 0, f = LU.Formulae.size(); i != f; ++i)
       GenerateScales(LU, LUIdx, LU.Formulae[i]);
   }
   for (size_t LUIdx = 0, NumUses = Uses.size(); LUIdx != NumUses; ++LUIdx) {
     LSRUse &LU = Uses[LUIdx];
     for (size_t i = 0, f = LU.Formulae.size(); i != f; ++i)
       GenerateTruncates(LU, LUIdx, LU.Formulae[i]);
+    for (size_t i = 0, f = LU.Formulae.size(); i != f; ++i)
+      GenerateZExts(LU, LUIdx, LU.Formulae[i]);
   }
 
   GenerateCrossUseConstantOffsets();
 
   DEBUG(dbgs() << "\n"
                   "After generating reuse formulae:\n";
         print_uses(dbgs()));
 }
Value *LSRInstance::Expand(const LSRFixup &LF,
   SmallVector<const SCEV *, 8> Ops;
 
   // Expand the BaseRegs portion.
   for (const SCEV *Reg : F.BaseRegs) {
     assert(!Reg->isZero() && "Zero allocated in a base register!");
 
     // If we're expanding for a post-inc user, make the post-inc adjustment.
     PostIncLoopSet &Loops = const_cast<PostIncLoopSet &>(LF.PostIncLoops);
-    Reg = TransformForPostIncUse(Denormalize, Reg,
-                                 LF.UserInst, LF.OperandValToReplace,
-                                 Loops, SE, DT);
+    const SCEV *ExtendedReg =
+        F.ZeroExtendBaseReg ? SE.getZeroExtendExpr(Reg, F.ZeroExtendType) : Reg;
 
-    Ops.push_back(SE.getUnknown(Rewriter.expandCodeFor(Reg, nullptr, IP)));
+    const SCEV *PostIncReg =
+        TransformForPostIncUse(Denormalize, ExtendedReg, LF.UserInst,
+                               LF.OperandValToReplace, Loops, SE, DT);
+    if (PostIncReg == ExtendedReg) {
+      Value *Expanded = Rewriter.expandCodeFor(Reg, nullptr, IP);
+      if (F.ZeroExtendBaseReg)
+        Expanded = new ZExtInst(Expanded, F.ZeroExtendType, "", IP);
+      Ops.push_back(SE.getUnknown(Expanded));
+    } else {
+      Ops.push_back(
+          SE.getUnknown(Rewriter.expandCodeFor(PostIncReg, nullptr, IP)));
     }
+  }
+
+  // Note on post-inc uses and zero extends -- since the no-wrap behavior for
+  // the post-inc SCEV can be different from the no-wrap behavior of the pre-inc
+  // SCEV, if a post-inc transform is required we do the zero extension on the
+  // pre-inc expression before doing the post-inc transform.
 
   // Expand the ScaledReg portion.
   Value *ICmpScaledV = nullptr;
   if (F.Scale != 0) {
     const SCEV *ScaledS = F.ScaledReg;
 
     // If we're expanding for a post-inc user, make the post-inc adjustment.
     PostIncLoopSet &Loops = const_cast<PostIncLoopSet &>(LF.PostIncLoops);
-    ScaledS = TransformForPostIncUse(Denormalize, ScaledS,
-                                     LF.UserInst, LF.OperandValToReplace,
-                                     Loops, SE, DT);
+    const SCEV *ExtendedScaleS =
+        F.ZeroExtendScaledReg ? SE.getZeroExtendExpr(ScaledS, F.ZeroExtendType)
+                              : ScaledS;
+    const SCEV *PostIncScaleS =
+        TransformForPostIncUse(Denormalize, ExtendedScaleS, LF.UserInst,
+                               LF.OperandValToReplace, Loops, SE, DT);
 
     if (LU.Kind == LSRUse::ICmpZero) {
       // Expand ScaleReg as if it was part of the base regs.
+      Value *Expanded = nullptr;
+      if (PostIncScaleS == ExtendedScaleS) {
+        Expanded = Rewriter.expandCodeFor(ScaledS, nullptr, IP);
+        if (F.ZeroExtendScaledReg)
+          Expanded = new ZExtInst(Expanded, F.ZeroExtendType, "", IP);
+      } else {
+        Expanded = Rewriter.expandCodeFor(PostIncScaleS, nullptr, IP);
+      }
+
       if (F.Scale == 1)
-        Ops.push_back(
-            SE.getUnknown(Rewriter.expandCodeFor(ScaledS, nullptr, IP)));
+        Ops.push_back(SE.getUnknown(Expanded));
       else {
         // An interesting way of "folding" with an icmp is to use a negated
         // scale, which we'll implement by inserting it into the other operand
         // of the icmp.
         assert(F.Scale == -1 &&
                "The only scale supported by ICmpZero uses is -1!");
-        ICmpScaledV = Rewriter.expandCodeFor(ScaledS, nullptr, IP);
+        ICmpScaledV = Expanded;
       }
     } else {
       // Otherwise just expand the scaled register and an explicit scale,
       // which is expected to be matched as part of the address.
 
       // Flush the operand list to suppress SCEVExpander hoisting address modes.
       // Unless the addressing mode will not be folded.
       if (!Ops.empty() && LU.Kind == LSRUse::Address &&
           isAMCompletelyFolded(TTI, LU, F)) {
         Value *FullV = Rewriter.expandCodeFor(SE.getAddExpr(Ops), Ty, IP);
         Ops.clear();
         Ops.push_back(SE.getUnknown(FullV));
       }
-      ScaledS = SE.getUnknown(Rewriter.expandCodeFor(ScaledS, nullptr, IP));
+
+      Value *Expanded = nullptr;
+      if (PostIncScaleS == ExtendedScaleS) {
+        Expanded = Rewriter.expandCodeFor(ScaledS, nullptr, IP);
+        if (F.ZeroExtendScaledReg)
+          Expanded = new ZExtInst(Expanded, F.ZeroExtendType, "", IP);
+      } else {
+        Expanded = Rewriter.expandCodeFor(PostIncScaleS, nullptr, IP);
+      }
+
+      ScaledS = SE.getUnknown(Expanded);
       if (F.Scale != 1)
         ScaledS =
             SE.getMulExpr(ScaledS, SE.getConstant(ScaledS->getType(), F.Scale));
       Ops.push_back(ScaledS);
     }
   }
 
   // Expand the GV portion.

llvm/trunk/test/Transforms/LoopStrengthReduce/zext-of-scale.ll

+; RUN: opt  < %s -S -loop-reduce | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+%struct = type { [8 x i8] }
+
+declare void @use_32(i32)
+declare void @use_64(i64)
+
+define void @f(i32 %tmp156, i32* %length_buf_1, i32* %length_buf_0, %struct* %b,
+                %struct* %c, %struct* %d, %struct* %e, i32* %length_buf_2,
+                i32 %tmp160) {
+; CHECK-LABEL: @f(
+entry:
+  %begin151 = getelementptr inbounds %struct, %struct* %b, i64 0, i32 0, i64 12
+  %tmp21 = bitcast i8* %begin151 to i32*
+  %begin157 = getelementptr inbounds %struct, %struct* %c, i64 0, i32 0, i64 16
+  %tmp23 = bitcast i8* %begin157 to double*
+  %begin163 = getelementptr inbounds %struct, %struct* %d, i64 0, i32 0, i64 16
+  %tmp25 = bitcast i8* %begin163 to double*
+  %length.i820 = load i32, i32* %length_buf_1, align 4, !range !0
+  %enter = icmp ne i32 %tmp156, -1
+  br i1 %enter, label %ok_146, label %block_81_2
+
+ok_146:
+  %var_13 = phi double [ %tmp186, %ok_161 ], [ 0.000000e+00, %entry ]
+  %var_17 = phi i32 [ %tmp187, %ok_161 ], [ %tmp156, %entry ]
+  %tmp174 = zext i32 %var_17 to i64
+  %tmp175 = icmp ult i32 %var_17, %length.i820
+  br i1 %tmp175, label %ok_152, label %block_81_2
+
+ok_152:
+  %tmp176 = getelementptr inbounds i32, i32* %tmp21, i64 %tmp174
+  %tmp177 = load i32, i32* %tmp176, align 4
+  %tmp178 = zext i32 %tmp177 to i64
+  %length.i836 = load i32, i32* %length_buf_2, align 4, !range !0
+  %tmp179 = icmp ult i32 %tmp177, %length.i836
+  br i1 %tmp179, label %ok_158, label %block_81_2
+
+ok_158:
+  %tmp180 = getelementptr inbounds double, double* %tmp23, i64 %tmp178
+  %tmp181 = load double, double* %tmp180, align 8
+  %length.i = load i32, i32* %length_buf_0, align 4, !range !0
+  %tmp182 = icmp slt i32 %var_17, %length.i
+  br i1 %tmp182, label %ok_161, label %block_81_2
+
+ok_161:
+; CHECK-LABEL: ok_161:
+; CHECK: add
+; CHECK-NOT: add
+  %tmp183 = getelementptr inbounds double, double* %tmp25, i64 %tmp174
+  %tmp184 = load double, double* %tmp183, align 8
+  %tmp185 = fmul double %tmp181, %tmp184
+  %tmp186 = fadd double %var_13, %tmp185
+  %tmp187 = add nsw i32 %var_17, 1
+  %tmp188 = icmp slt i32 %tmp187, %tmp160
+; CHECK: br
+  br i1 %tmp188, label %ok_146, label %block_81
+
+block_81:
+  call void @use_64(i64 %tmp174)  ;; pre-inc use
+  call void @use_32(i32 %tmp187)  ;; post-inc use
+  ret void
+
+block_81_2:
+  ret void
+}
+
+!0 = !{i32 0, i32 2147483647}