1 files changed, 82 insertions, 6 deletions

diff --git a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
index 93a3b9281d9..b74eaf0daf5 100644
--- a/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
+++ b/llvm/lib/Target/X86/X86ISelDAGToDAG.cpp
@@ -88,6 +88,11 @@ namespace {
              IndexReg.getNode() != nullptr || Base_Reg.getNode() != nullptr;
     }
 
+    bool hasComplexAddressingMode() const {
+      return Disp && IndexReg.getNode() != nullptr &&
+             Base_Reg.getNode() != nullptr;
+    }
+
     /// Return true if this addressing mode is already RIP-relative.
     bool isRIPRelative() const {
       if (BaseType != RegBase) return false;
@@ -97,6 +102,10 @@ namespace {
       return false;
     }
 
+    bool isLegalScale() const {
+      return (Scale == 1 || Scale == 2 || Scale == 4 || Scale == 8);
+    }
+
     void setBaseReg(SDValue Reg) {
       BaseType = RegBase;
       Base_Reg = Reg;
@@ -162,10 +171,13 @@ namespace {
     /// If true, selector should try to optimize for minimum code size.
     bool OptForMinSize;
 
+    /// If true, selector should try to aggresively fold operands into AM.
+    bool OptForAggressingFolding;
+
   public:
     explicit X86DAGToDAGISel(X86TargetMachine &tm, CodeGenOpt::Level OptLevel)
         : SelectionDAGISel(tm, OptLevel), OptForSize(false),
-          OptForMinSize(false) {}
+          OptForMinSize(false), OptForAggressingFolding(false) {}
 
     StringRef getPassName() const override {
       return "X86 DAG->DAG Instruction Selection";
@@ -184,6 +196,12 @@ namespace {
 
     void PreprocessISelDAG() override;
 
+    void setAggressiveOperandFolding(bool val) {
+      OptForAggressingFolding = val;
+    }
+
+    bool getAggressiveOperandFolding() { return OptForAggressingFolding; }
+
 // Include the pieces autogenerated from the target description.
 #include "X86GenDAGISel.inc"
 
@@ -198,6 +216,7 @@ namespace {
     bool matchAdd(SDValue N, X86ISelAddressMode &AM, unsigned Depth);
     bool matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
                                  unsigned Depth);
+    bool matchAddressLEA(SDValue N, X86ISelAddressMode &AM);
     bool matchAddressBase(SDValue N, X86ISelAddressMode &AM);
     bool selectAddr(SDNode *Parent, SDValue N, SDValue &Base,
                     SDValue &Scale, SDValue &Index, SDValue &Disp,
@@ -447,6 +466,20 @@ namespace {
 
     bool isMaskZeroExtended(SDNode *N) const;
   };
+
+  class X86AggressiveOperandFolding {
+  public:
+    explicit X86AggressiveOperandFolding(X86DAGToDAGISel &ISel, bool val)
+        : Selector(&ISel) {
+      Selector->setAggressiveOperandFolding(val);
+    }
+    ~X86AggressiveOperandFolding() {
+      Selector->setAggressiveOperandFolding(false);
+    }
+
+  private:
+    X86DAGToDAGISel *Selector;
+  };
 }
 
 
@@ -1194,7 +1227,7 @@ static bool foldMaskAndShiftToScale(SelectionDAG &DAG, SDValue N,
   AM.IndexReg = NewSRL;
   return false;
 }
-
+   
 bool X86DAGToDAGISel::matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
                                               unsigned Depth) {
   SDLoc dl(N);
@@ -1202,8 +1235,14 @@ bool X86DAGToDAGISel::matchAddressRecursively(SDValue N, X86ISelAddressMode &AM,
       dbgs() << "MatchAddress: ";
       AM.dump();
     });
-  // Limit recursion.
-  if (Depth > 5)
+ 
+  // Limit recursion. For aggressive operand folding recurse 
+  // till depth 8 which is the maximum legal scale value.
+  auto getMaxOperandFoldingDepth = [&] () -> unsigned int {
+      return getAggressiveOperandFolding() ? 8 : 5;
+  };
+
+  if (Depth > getMaxOperandFoldingDepth())
     return matchAddressBase(N, AM);
 
   // If this is already a %rip relative address, we can only merge immediates
@@ -1494,6 +1533,20 @@ bool X86DAGToDAGISel::matchAddressBase(SDValue N, X86ISelAddressMode &AM) {
       return false;
     }
 
+    if (OptLevel != CodeGenOpt::None && getAggressiveOperandFolding() &&
+        AM.BaseType == X86ISelAddressMode::RegBase) {
+      if (AM.Base_Reg == N) {
+        SDValue Base_Reg = AM.Base_Reg;
+        AM.Base_Reg = AM.IndexReg;
+        AM.IndexReg = Base_Reg;
+        AM.Scale++;
+        return false;
+      } else if (AM.IndexReg == N) {
+        AM.Scale++;
+        return false;
+      }
+    }
+
     // Otherwise, we cannot select it.
     return true;
   }
@@ -1729,7 +1782,7 @@ bool X86DAGToDAGISel::selectLEA64_32Addr(SDValue N, SDValue &Base,
                                          SDValue &Disp, SDValue &Segment) {
   // Save the debug loc before calling selectLEAAddr, in case it invalidates N.
   SDLoc DL(N);
-
+   
   if (!selectLEAAddr(N, Base, Scale, Index, Disp, Segment))
     return false;
 
@@ -1764,6 +1817,29 @@ bool X86DAGToDAGISel::selectLEA64_32Addr(SDValue N, SDValue &Base,
   return true;
 }
 
+bool X86DAGToDAGISel::matchAddressLEA(SDValue N, X86ISelAddressMode &AM) {
+  // Avoid enabling aggressive operand folding when node N is a part of loop.
+  X86AggressiveOperandFolding Enable(*this, !CurDAG->IsDAGPartOfLoop);
+
+  bool matchRes = matchAddress(N, AM);
+
+  // Check for legality of scale when recursion unwinds back to the top.
+  if (!matchRes) {
+    if (!AM.isLegalScale())
+      return true;
+
+    // Avoid creating costly complex LEAs having scale less than 2
+    // within loop.
+    if(CurDAG->IsDAGPartOfLoop && Subtarget->slow3OpsLEA() &&
+        AM.Scale <= 2 && AM.hasComplexAddressingMode() &&
+         (!AM.hasSymbolicDisplacement() && N.getOpcode() < ISD::BUILTIN_OP_END))
+     return true;
+  }
+
+  return matchRes;
+}
+
+
 /// Calls SelectAddr and determines if the maximal addressing
 /// mode it matches can be cost effectively emitted as an LEA instruction.
 bool X86DAGToDAGISel::selectLEAAddr(SDValue N,
@@ -1781,7 +1857,7 @@ bool X86DAGToDAGISel::selectLEAAddr(SDValue N,
   SDValue Copy = AM.Segment;
   SDValue T = CurDAG->getRegister(0, MVT::i32);
   AM.Segment = T;
-  if (matchAddress(N, AM))
+  if (matchAddressLEA(N, AM))
     return false;
   assert (T == AM.Segment);
   AM.Segment = Copy;