[x86] X86ISelLowering zext(add_nuw(x, C)) --> add(zext(x), C_zext)

[apilipenko]

Summary

Currently X86ISelLowering has a similar transformation for sexts:
sext(add_nsw(x, C)) --> add(sext(x), C_sext)
See D13757 for motivation.

In this change I extend this code to handle zexts as well.

[spatel]

LGTM. See inline comment for one nit. Also, unless there's some reason to break this regression test off by itself, I'd prefer to add this test directly to the related test file (and rename the file so it's not inaccurate). Keeping the tests next to each other for context seems better to me.

lib/Target/X86/X86ISelLowering.cpp

30492

This should be named 'NewExt' now for consistency.

lib/Target/X86/X86ISelLowering.cpp

                                   N00, N1);
       return DAG.getNode(ISD::SIGN_EXTEND, dl, MVT::v4i64, Tmp);
     }
   }
   return SDValue();
 }
 
 /// sext(add_nsw(x, C)) --> add(sext(x), C_sext)
-/// Promoting a sign extension ahead of an 'add nsw' exposes opportunities
-/// to combine math ops, use an LEA, or use a complex addressing mode. This can
-/// eliminate extend, add, and shift instructions.
-static SDValue promoteSextBeforeAddNSW(SDNode *Sext, SelectionDAG &DAG,
-                                       const X86Subtarget &Subtarget) {
+/// zext(add_nuw(x, C)) --> add(zext(x), C_zext)
+/// Promoting a sign/zero extension ahead of a no overflow 'add' exposes
+/// opportunities to combine math ops, use an LEA, or use a complex addressing
+/// mode. This can eliminate extend, add, and shift instructions.
+static SDValue promoteExtBeforeAdd(SDNode *Ext, SelectionDAG &DAG,
+                                   const X86Subtarget &Subtarget) {
+  if (Ext->getOpcode() != ISD::SIGN_EXTEND &&
+      Ext->getOpcode() != ISD::ZERO_EXTEND)
+    return SDValue();
+
   // TODO: This should be valid for other integer types.
-  EVT VT = Sext->getValueType(0);
+  EVT VT = Ext->getValueType(0);
   if (VT != MVT::i64)
     return SDValue();
 
-  // We need an 'add nsw' feeding into the 'sext'.
-  SDValue Add = Sext->getOperand(0);
-  if (Add.getOpcode() != ISD::ADD || !Add->getFlags()->hasNoSignedWrap())
+  SDValue Add = Ext->getOperand(0);
+  if (Add.getOpcode() != ISD::ADD)
+    return SDValue();
+
+  bool Sext = Ext->getOpcode() == ISD::SIGN_EXTEND;
+  bool NSW = Add->getFlags()->hasNoSignedWrap();
+  bool NUW = Add->getFlags()->hasNoUnsignedWrap();
+
+  // We need an 'add nsw' feeding into the 'sext' or 'add nuw' feeding
+  // into the 'zext'
+  if ((Sext && !NSW) || (!Sext && !NUW))
     return SDValue();
 
   // Having a constant operand to the 'add' ensures that we are not increasing
   // the instruction count because the constant is extended for free below.
   // A constant operand can also become the displacement field of an LEA.
   auto *AddOp1 = dyn_cast<ConstantSDNode>(Add.getOperand(1));
   if (!AddOp1)
     return SDValue();
 
   // Don't make the 'add' bigger if there's no hope of combining it with some
   // other 'add' or 'shl' instruction.
   // TODO: It may be profitable to generate simpler LEA instructions in place
   // of single 'add' instructions, but the cost model for selecting an LEA
   // currently has a high threshold.
   bool HasLEAPotential = false;
-  for (auto *User : Sext->uses()) {
+  for (auto *User : Ext->uses()) {
     if (User->getOpcode() == ISD::ADD || User->getOpcode() == ISD::SHL) {
       HasLEAPotential = true;
       break;
     }
   }
   if (!HasLEAPotential)
     return SDValue();
 
-  // Everything looks good, so pull the 'sext' ahead of the 'add'.
-  int64_t AddConstant = AddOp1->getSExtValue();
+  // Everything looks good, so pull the '{s|z}ext' ahead of the 'add'.
+  int64_t AddConstant = Sext ? AddOp1->getSExtValue() : AddOp1->getZExtValue();
   SDValue AddOp0 = Add.getOperand(0);
-  SDValue NewSext = DAG.getNode(ISD::SIGN_EXTEND, SDLoc(Sext), VT, AddOp0);
+  SDValue NewSext = DAG.getNode(Ext->getOpcode(), SDLoc(Ext), VT, AddOp0);
   SDValue NewConstant = DAG.getConstant(AddConstant, SDLoc(Add), VT);
 
   // The wider add is guaranteed to not wrap because both operands are
   // sign-extended.
   SDNodeFlags Flags;
-  Flags.setNoSignedWrap(true);
+  Flags.setNoSignedWrap(NSW);
+  Flags.setNoUnsignedWrap(NUW);
   return DAG.getNode(ISD::ADD, SDLoc(Add), VT, NewSext, NewConstant, &Flags);
 }
 
 /// (i8,i32 {s/z}ext ({s/u}divrem (i8 x, i8 y)) ->
 /// (i8,i32 ({s/u}divrem_sext_hreg (i8 x, i8 y)
 /// This exposes the {s/z}ext to the sdivrem lowering, so that it directly
 /// extends from AH (which we otherwise need to do contortions to access).
 static SDValue getDivRem8(SDNode *N, SelectionDAG &DAG) {

 
   if (SDValue V = combineToExtendVectorInReg(N, DAG, DCI, Subtarget))
     return V;
 
   if (Subtarget.hasAVX() && VT.is256BitVector())
     if (SDValue R = WidenMaskArithmetic(N, DAG, DCI, Subtarget))
       return R;
 
-  if (SDValue NewAdd = promoteSextBeforeAddNSW(N, DAG, Subtarget))
+  if (SDValue NewAdd = promoteExtBeforeAdd(N, DAG, Subtarget))
     return NewAdd;
 
   return SDValue();
 }
 
 static SDValue combineFMA(SDNode *N, SelectionDAG &DAG,
                           const X86Subtarget &Subtarget) {
   SDLoc dl(N);

 
   if (VT.is256BitVector())
     if (SDValue R = WidenMaskArithmetic(N, DAG, DCI, Subtarget))
       return R;
 
   if (SDValue DivRem8 = getDivRem8(N, DAG))
     return DivRem8;
 
+  if (SDValue NewAdd = promoteExtBeforeAdd(N, DAG, Subtarget))
+    return NewAdd;
+
   return SDValue();
 }
 
 /// Optimize x == -y --> x+y == 0
 ///          x != -y --> x+y != 0
 static SDValue combineSetCC(SDNode *N, SelectionDAG &DAG,
                             const X86Subtarget &Subtarget) {
   ISD::CondCode CC = cast<CondCodeSDNode>(N->getOperand(2))->get();

test/CodeGen/X86/add-nuw-zext.ll

+; RUN: llc < %s -mtriple=x86_64-unknown-unknown | FileCheck %s
+
+; This is the same test as @PR20134 from add-nsw-sext.ll but sign extension is
+; replaced with zero extension. We must be able to fold additions into complex
+; addressing mode in this case as well.
+
+define void @PR20134(i32* %a, i32 %i) {
+; CHECK: # BB#0:
+; CHECK-NEXT: movl %esi, %eax
+; CHECK-NEXT: movl 4(%rdi,%rax,4), %ecx
+; CHECK-NEXT: addl 8(%rdi,%rax,4), %ecx
+; CHECK-NEXT: movl %ecx, (%rdi,%rax,4)
+; CHECK-NEXT: retq
+
+  %add1 = add nuw i32 %i, 1
+  %idx1 = zext i32 %add1 to i64
+  %gep1 = getelementptr i32, i32* %a, i64 %idx1
+  %load1 = load i32, i32* %gep1, align 4
+
+  %add2 = add nuw i32 %i, 2
+  %idx2 = zext i32 %add2 to i64
+  %gep2 = getelementptr i32, i32* %a, i64 %idx2
+  %load2 = load i32, i32* %gep2, align 4
+
+  %add3 = add i32 %load1, %load2
+  %idx3 = zext i32 %i to i64
+  %gep3 = getelementptr i32, i32* %a, i64 %idx3
+  store i32 %add3, i32* %gep3, align 4
+  ret void
+}