[X86] promoteExtBeforeAdd - add support for or/xor 'addlike' patterns

Fold zext(addlike(x, C)) --> add(zext(x), C_zext) if its likely to help us create LEA instructions

Addresses some regressions exposed by D155472

Author: RKSimon

llvm/lib/Target/X86/X86ISelLowering.cpp

@@ -52131,11 +52131,12 @@ static SDValue combineSignExtendInReg(SDNode *N, SelectionDAG &DAG,
   return SDValue();
 }
 
-/// sext(add_nsw(x, C)) --> add(sext(x), C_sext)
-/// 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.
+/// sext(add_nsw(x, C)) --> add_nsw(sext(x), C_sext)
+/// zext(add_nuw(x, C)) --> add_nuw(zext(x), C_zext)
+/// zext(addlike(x, C)) --> add(zext(x), C_zext)
+/// Promoting a sign/zero extension ahead of a no overflow 'add' or 'addlike'
+/// 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 &&
@@ -52147,17 +52148,19 @@ static SDValue promoteExtBeforeAdd(SDNode *Ext, SelectionDAG &DAG,
   if (VT != MVT::i64)
     return SDValue();
 
-  SDValue Add = Ext->getOperand(0);
-  if (Add.getOpcode() != ISD::ADD)
-    return SDValue();
-
+  bool NSW = false, NUW = false;
   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))
+  SDValue Add = Ext->getOperand(0);
+  unsigned AddOpc = Add->getOpcode();
+  if (AddOpc == ISD::ADD) {
+    NSW = Add->getFlags().hasNoSignedWrap();
+    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();
+  } else if (!(!Sext && DAG.isADDLike(Add)))
     return SDValue();
 
   // Having a constant operand to the 'add' ensures that we are not increasing
@@ -52193,7 +52196,7 @@ static SDValue promoteExtBeforeAdd(SDNode *Ext, SelectionDAG &DAG,
   SDNodeFlags Flags;
   Flags.setNoSignedWrap(NSW);
   Flags.setNoUnsignedWrap(NUW);
-  return DAG.getNode(ISD::ADD, SDLoc(Add), VT, NewExt, NewConstant, Flags);
+  return DAG.getNode(AddOpc, SDLoc(Add), VT, NewExt, NewConstant, Flags);
 }
 
 // If we face {ANY,SIGN,ZERO}_EXTEND that is applied to a CMOV with constant

llvm/test/CodeGen/X86/lea-2.ll

@@ -26,7 +26,7 @@ define i32 @test1(i32 %A, i32 %B) {
   ret i32 %t4
 }
 
-; TODO: The addlike OR instruction should fold into the LEA.
+; The addlike OR instruction should fold into the LEA.
 
 define i64 @test2(i32 %a0, i64 %a1) {
 ; X32-LABEL: test2:
@@ -44,8 +44,7 @@ define i64 @test2(i32 %a0, i64 %a1) {
 ; X64:       # %bb.0:
 ; X64-NEXT:    # kill: def $edi killed $edi def $rdi
 ; X64-NEXT:    andl $-8, %edi
-; X64-NEXT:    orl $2, %edi
-; X64-NEXT:    leaq (%rsi,%rdi,2), %rax
+; X64-NEXT:    leaq 4(%rsi,%rdi,2), %rax
 ; X64-NEXT:    retq
   %x1 = and i32 %a0, -8
   %x2 = or i32 %x1, 2

llvm/test/CodeGen/X86/select_const.ll

@@ -626,8 +626,8 @@ define i64 @select_pow2_diff_neg_invert(i1 zeroext %cond) {
 ;
 ; X64-LABEL: select_pow2_diff_neg_invert:
 ; X64:       # %bb.0:
-; X64-NEXT:    xorb $1, %dil
-; X64-NEXT:    movzbl %dil, %eax
+; X64-NEXT:    movl %edi, %eax
+; X64-NEXT:    xorq $1, %rax
 ; X64-NEXT:    shlq $7, %rax
 ; X64-NEXT:    addq $-99, %rax
 ; X64-NEXT:    retq