[X86] matchIndexRecursively - add zext(add/addlike(x,c)) -> index: zext(x), disp + zext(c) pattern handling

More restricted alternative to a8cef6b

Author: RKSimon

llvm/lib/Target/X86/X86ISelDAGToDAG.cpp

@@ -2280,6 +2280,35 @@ SDValue X86DAGToDAGISel::matchIndexRecursively(SDValue N,
     }
   }
 
+  // index: zext(add_nuw(x,c)) -> index: zext(x), disp + zext(c)
+  // index: zext(addlike(x,c)) -> index: zext(x), disp + zext(c)
+  // TODO: call matchIndexRecursively(AddSrc) if we won't corrupt sext?
+  if (Opc == ISD::ZERO_EXTEND && !VT.isVector() && N.hasOneUse()) {
+    SDValue Src = N.getOperand(0);
+    unsigned SrcOpc = Src.getOpcode();
+    if (((SrcOpc == ISD::ADD && Src->getFlags().hasNoUnsignedWrap()) ||
+         CurDAG->isADDLike(Src)) &&
+        Src.hasOneUse()) {
+      if (CurDAG->isBaseWithConstantOffset(Src)) {
+        SDValue AddSrc = Src.getOperand(0);
+        auto *AddVal = cast<ConstantSDNode>(Src.getOperand(1));
+        uint64_t Offset = (uint64_t)AddVal->getZExtValue();
+        if (!foldOffsetIntoAddress(Offset * AM.Scale, AM)) {
+          SDLoc DL(N);
+          SDValue ExtSrc = CurDAG->getNode(Opc, DL, VT, AddSrc);
+          SDValue ExtVal = CurDAG->getConstant(Offset, DL, VT);
+          SDValue ExtAdd = CurDAG->getNode(SrcOpc, DL, VT, ExtSrc, ExtVal);
+          insertDAGNode(*CurDAG, N, ExtSrc);
+          insertDAGNode(*CurDAG, N, ExtVal);
+          insertDAGNode(*CurDAG, N, ExtAdd);
+          CurDAG->ReplaceAllUsesWith(N, ExtAdd);
+          CurDAG->RemoveDeadNode(N.getNode());
+          return ExtSrc;
+        }
+      }
+    }
+  }
+
   // TODO: Handle extensions, shifted masks etc.
   return N;
 }

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