1 files changed, 78 insertions, 36 deletions

diff --git a/llvm/lib/Target/X86/X86OptimizeLEAs.cpp b/llvm/lib/Target/X86/X86OptimizeLEAs.cpp
index def49dfbe5c..1a3a0951a74 100644
--- a/llvm/lib/Target/X86/X86OptimizeLEAs.cpp
+++ b/llvm/lib/Target/X86/X86OptimizeLEAs.cpp
@@ -24,6 +24,7 @@
 #include "llvm/CodeGen/LiveVariables.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineOperand.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/IR/Function.h"
@@ -53,6 +54,17 @@ static inline MemOpKey getMemOpKey(const MachineInstr &MI, unsigned N);
 static inline bool isIdenticalOp(const MachineOperand &MO1,
                                  const MachineOperand &MO2);
 
+/// \brief Returns true if two address displacement operands are of the same
+/// type and use the same symbol/index/address regardless of the offset.
+static bool isSimilarDispOp(const MachineOperand &MO1,
+                            const MachineOperand &MO2);
+
+/// \brief Returns true if the type of \p MO is valid for address displacement
+/// operand. According to X86DAGToDAGISel::getAddressOperands allowed types are:
+/// MO_Immediate, MO_ConstantPoolIndex, MO_JumpTableIndex, MO_ExternalSymbol,
+/// MO_GlobalAddress, MO_BlockAddress or MO_MCSymbol.
+static inline bool isValidDispOp(const MachineOperand &MO);
+
 /// \brief Returns true if the instruction is LEA.
 static inline bool isLEA(const MachineInstr &MI);
 
@@ -75,19 +87,11 @@ public:
       if (!isIdenticalOp(*Operands[i], *Other.Operands[i]))
         return false;
 
-    assert((Disp->isImm() || Disp->isGlobal()) &&
-           (Other.Disp->isImm() || Other.Disp->isGlobal()) &&
-           "Address displacement operand is always an immediate or a global");
-
-    // Addresses' displacements must be either immediates or the same global.
-    // Immediates' and offsets' values don't matter for the operator since the
-    // difference will be taken care of during instruction substitution.
-    if ((Disp->isImm() && Other.Disp->isImm()) ||
-        (Disp->isGlobal() && Other.Disp->isGlobal() &&
-         Disp->getGlobal() == Other.Disp->getGlobal()))
-      return true;
-
-    return false;
+    // Addresses' displacements don't have to be exactly the same. It only
+    // matters that they use the same symbol/index/address. Immediates' or
+    // offsets' differences will be taken care of during instruction
+    // substitution.
+    return isSimilarDispOp(*Disp, *Other.Disp);
   }
 
   // Address' base, scale, index and segment operands.
@@ -126,10 +130,30 @@ template <> struct DenseMapInfo<MemOpKey> {
 
     // If the address displacement is an immediate, it should not affect the
     // hash so that memory operands which differ only be immediate displacement
-    // would have the same hash. If the address displacement is a global, we
-    // should reflect this global in the hash.
-    if (Val.Disp->isGlobal())
+    // would have the same hash. If the address displacement is something else,
+    // we should reflect symbol/index/address in the hash.
+    switch (Val.Disp->getType()) {
+    case MachineOperand::MO_Immediate:
+      break;
+    case MachineOperand::MO_ConstantPoolIndex:
+    case MachineOperand::MO_JumpTableIndex:
+      Hash = hash_combine(Hash, Val.Disp->getIndex());
+      break;
+    case MachineOperand::MO_ExternalSymbol:
+      Hash = hash_combine(Hash, Val.Disp->getSymbolName());
+      break;
+    case MachineOperand::MO_GlobalAddress:
       Hash = hash_combine(Hash, Val.Disp->getGlobal());
+      break;
+    case MachineOperand::MO_BlockAddress:
+      Hash = hash_combine(Hash, Val.Disp->getBlockAddress());
+      break;
+    case MachineOperand::MO_MCSymbol:
+      Hash = hash_combine(Hash, Val.Disp->getMCSymbol());
+      break;
+    default:
+      llvm_unreachable("Invalid address displacement operand");
+    }
 
     return (unsigned)Hash;
   }
@@ -163,6 +187,28 @@ static inline bool isIdenticalOp(const MachineOperand &MO1,
           !TargetRegisterInfo::isPhysicalRegister(MO1.getReg()));
 }
 
+static bool isSimilarDispOp(const MachineOperand &MO1,
+                            const MachineOperand &MO2) {
+  assert(isValidDispOp(MO1) && isValidDispOp(MO2) &&
+         "Address displacement operand is not valid");
+  return (MO1.isImm() && MO2.isImm()) ||
+         (MO1.isCPI() && MO2.isCPI() && MO1.getIndex() == MO2.getIndex()) ||
+         (MO1.isJTI() && MO2.isJTI() && MO1.getIndex() == MO2.getIndex()) ||
+         (MO1.isSymbol() && MO2.isSymbol() &&
+          MO1.getSymbolName() == MO2.getSymbolName()) ||
+         (MO1.isGlobal() && MO2.isGlobal() &&
+          MO1.getGlobal() == MO2.getGlobal()) ||
+         (MO1.isBlockAddress() && MO2.isBlockAddress() &&
+          MO1.getBlockAddress() == MO2.getBlockAddress()) ||
+         (MO1.isMCSymbol() && MO2.isMCSymbol() &&
+          MO1.getMCSymbol() == MO2.getMCSymbol());
+}
+
+static inline bool isValidDispOp(const MachineOperand &MO) {
+  return MO.isImm() || MO.isCPI() || MO.isJTI() || MO.isSymbol() ||
+         MO.isGlobal() || MO.isBlockAddress() || MO.isMCSymbol();
+}
+
 static inline bool isLEA(const MachineInstr &MI) {
   unsigned Opcode = MI.getOpcode();
   return Opcode == X86::LEA16r || Opcode == X86::LEA32r ||
@@ -317,16 +363,19 @@ bool OptimizeLEAPass::chooseBestLEA(const SmallVectorImpl<MachineInstr *> &List,
 int64_t OptimizeLEAPass::getAddrDispShift(const MachineInstr &MI1, unsigned N1,
                                           const MachineInstr &MI2,
                                           unsigned N2) const {
-  // Address displacement operands may differ by a constant.
   const MachineOperand &Op1 = MI1.getOperand(N1 + X86::AddrDisp);
   const MachineOperand &Op2 = MI2.getOperand(N2 + X86::AddrDisp);
-  if (Op1.isImm() && Op2.isImm())
-    return Op1.getImm() - Op2.getImm();
-  else if (Op1.isGlobal() && Op2.isGlobal() &&
-           Op1.getGlobal() == Op2.getGlobal())
-    return Op1.getOffset() - Op2.getOffset();
-  else
-    llvm_unreachable("Invalid address displacement operand");
+
+  assert(isSimilarDispOp(Op1, Op2) &&
+         "Address displacement operands are not compatible");
+
+  // After the assert above we can be sure that both operands are of the same
+  // valid type and use the same symbol/index/address, thus displacement shift
+  // calculation is rather simple.
+  if (Op1.isJTI())
+    return 0;
+  return Op1.isImm() ? Op1.getImm() - Op2.getImm()
+                     : Op1.getOffset() - Op2.getOffset();
 }
 
 // Check that the Last LEA can be replaced by the First LEA. To be so,
@@ -434,11 +483,6 @@ bool OptimizeLEAPass::removeRedundantAddrCalc(MemOpMap &LEAs) {
 
     MemOpNo += X86II::getOperandBias(Desc);
 
-    // Address displacement must be an immediate or a global.
-    MachineOperand &Disp = MI.getOperand(MemOpNo + X86::AddrDisp);
-    if (!Disp.isImm() && !Disp.isGlobal())
-      continue;
-
     // Get the best LEA instruction to replace address calculation.
     MachineInstr *DefMI;
     int64_t AddrDispShift;
@@ -536,13 +580,11 @@ bool OptimizeLEAPass::removeRedundantLEAs(MemOpMap &LEAs) {
           MO.setReg(First.getOperand(0).getReg());
 
           // Update address disp.
-          MachineOperand *Op = &MI.getOperand(MemOpNo + X86::AddrDisp);
-          if (Op->isImm())
-            Op->setImm(Op->getImm() + AddrDispShift);
-          else if (Op->isGlobal())
-            Op->setOffset(Op->getOffset() + AddrDispShift);
-          else
-            llvm_unreachable("Invalid address displacement operand");
+          MachineOperand &Op = MI.getOperand(MemOpNo + X86::AddrDisp);
+          if (Op.isImm())
+            Op.setImm(Op.getImm() + AddrDispShift);
+          else if (!Op.isJTI())
+            Op.setOffset(Op.getOffset() + AddrDispShift);
         }
 
         // Since we can possibly extend register lifetime, clear kill flags.