[X86] Update X86::getConstantFromPool to take base OperandNo instead of Displacement MachineOperand

This allows us to check the entire constant address calculation, and ensure we're not performing any runtime address math into the constant pool (noticed in an upcoming patch).

4 files changed, 22 insertions, 44 deletions

diff --git a/llvm/lib/Target/X86/X86FixupVectorConstants.cpp b/llvm/lib/Target/X86/X86FixupVectorConstants.cpp
index 5ae41dc..9da8d73 100644
--- a/llvm/lib/Target/X86/X86FixupVectorConstants.cpp
+++ b/llvm/lib/Target/X86/X86FixupVectorConstants.cpp
@@ -240,8 +240,7 @@ bool X86FixupVectorConstantsPass::processInstruction(MachineFunction &MF,
     assert(MI.getNumOperands() >= (OperandNo + X86::AddrNumOperands) &&
            "Unexpected number of operands!");
 
-    MachineOperand &CstOp = MI.getOperand(OperandNo + X86::AddrDisp);
-    if (auto *C = X86::getConstantFromPool(MI, CstOp)) {
+    if (auto *C = X86::getConstantFromPool(MI, OperandNo)) {
       // Attempt to detect a suitable splat from increasing splat widths.
       std::pair<unsigned, unsigned> Broadcasts[] = {
           {8, OpBcst8},   {16, OpBcst16},   {32, OpBcst32},
@@ -255,7 +254,7 @@ bool X86FixupVectorConstantsPass::processInstruction(MachineFunction &MF,
             unsigned NewCPI =
                 CP->getConstantPoolIndex(NewCst, Align(BitWidth / 8));
             MI.setDesc(TII->get(OpBcst));
-            CstOp.setIndex(NewCPI);
+            MI.getOperand(OperandNo + X86::AddrDisp).setIndex(NewCPI);
             return true;
           }
         }
diff --git a/llvm/lib/Target/X86/X86InstrInfo.cpp b/llvm/lib/Target/X86/X86InstrInfo.cpp
index 36022ef..6082a56 100644
--- a/llvm/lib/Target/X86/X86InstrInfo.cpp
+++ b/llvm/lib/Target/X86/X86InstrInfo.cpp
@@ -3515,13 +3515,21 @@ int X86::getFirstAddrOperandIdx(const MachineInstr &MI) {
 }
 
 const Constant *X86::getConstantFromPool(const MachineInstr &MI,
-                                         const MachineOperand &Op) {
-  if (!Op.isCPI() || Op.getOffset() != 0)
+                                         unsigned OpNo) {
+  assert(MI.getNumOperands() >= (OpNo + X86::AddrNumOperands) &&
+         "Unexpected number of operands!");
+
+  const MachineOperand &Index = MI.getOperand(OpNo + X86::AddrIndexReg);
+  if (!Index.isReg() || Index.getReg() != X86::NoRegister)
+    return nullptr;
+
+  const MachineOperand &Disp = MI.getOperand(OpNo + X86::AddrDisp);
+  if (!Disp.isCPI() || Disp.getOffset() != 0)
     return nullptr;
 
   ArrayRef<MachineConstantPoolEntry> Constants =
       MI.getParent()->getParent()->getConstantPool()->getConstants();
-  const MachineConstantPoolEntry &ConstantEntry = Constants[Op.getIndex()];
+  const MachineConstantPoolEntry &ConstantEntry = Constants[Disp.getIndex()];
 
   // Bail if this is a machine constant pool entry, we won't be able to dig out
   // anything useful.
diff --git a/llvm/lib/Target/X86/X86InstrInfo.h b/llvm/lib/Target/X86/X86InstrInfo.h
index 407bb87..939bc7da 100644
--- a/llvm/lib/Target/X86/X86InstrInfo.h
+++ b/llvm/lib/Target/X86/X86InstrInfo.h
@@ -87,8 +87,7 @@ bool isX87Instruction(MachineInstr &MI);
 int getFirstAddrOperandIdx(const MachineInstr &MI);
 
 /// Find any constant pool entry associated with a specific instruction operand.
-const Constant *getConstantFromPool(const MachineInstr &MI,
-                                    const MachineOperand &Op);
+const Constant *getConstantFromPool(const MachineInstr &MI, unsigned OpNo);
 
 } // namespace X86
 
diff --git a/llvm/lib/Target/X86/X86MCInstLower.cpp b/llvm/lib/Target/X86/X86MCInstLower.cpp
index c92712d..1c251e7 100644
--- a/llvm/lib/Target/X86/X86MCInstLower.cpp
+++ b/llvm/lib/Target/X86/X86MCInstLower.cpp
@@ -1539,16 +1539,12 @@ static void printConstant(const Constant *COp, unsigned BitWidth,
 static void printZeroUpperMove(const MachineInstr *MI, MCStreamer &OutStreamer,
                                int SclWidth, int VecWidth,
                                const char *ShuffleComment) {
-  assert(MI->getNumOperands() >= (1 + X86::AddrNumOperands) &&
-         "Unexpected number of operands!");
-
   std::string Comment;
   raw_string_ostream CS(Comment);
   const MachineOperand &DstOp = MI->getOperand(0);
   CS << X86ATTInstPrinter::getRegisterName(DstOp.getReg()) << " = ";
 
-  if (auto *C =
-          X86::getConstantFromPool(*MI, MI->getOperand(1 + X86::AddrDisp))) {
+  if (auto *C = X86::getConstantFromPool(*MI, 1)) {
     int CstEltSize = C->getType()->getScalarSizeInBits();
     if (SclWidth == CstEltSize) {
       if (auto *CI = dyn_cast<ConstantInt>(C)) {
@@ -1584,10 +1580,7 @@ static void printZeroUpperMove(const MachineInstr *MI, MCStreamer &OutStreamer,
 
 static void printLaneBroadcast(const MachineInstr *MI, MCStreamer &OutStreamer,
                                int NumLanes, int BitWidth) {
-  assert(MI->getNumOperands() >= (1 + X86::AddrNumOperands) &&
-         "Unexpected number of operands!");
-  if (auto *C =
-          X86::getConstantFromPool(*MI, MI->getOperand(1 + X86::AddrDisp))) {
+  if (auto *C = X86::getConstantFromPool(*MI, 1)) {
     int CstEltSize = C->getType()->getScalarSizeInBits();
 
     std::string Comment;
@@ -1637,10 +1630,7 @@ static void printLaneBroadcast(const MachineInstr *MI, MCStreamer &OutStreamer,
 static void printElementBroadcast(const MachineInstr *MI,
                                   MCStreamer &OutStreamer, int NumElts,
                                   int EltBits) {
-  assert(MI->getNumOperands() >= (1 + X86::AddrNumOperands) &&
-         "Unexpected number of operands!");
-  if (auto *C =
-          X86::getConstantFromPool(*MI, MI->getOperand(1 + X86::AddrDisp))) {
+  if (auto *C = X86::getConstantFromPool(*MI, 1)) {
     std::string Comment;
     raw_string_ostream CS(Comment);
     const MachineOperand &DstOp = MI->getOperand(0);
@@ -1771,13 +1761,8 @@ static void addConstantComments(const MachineInstr *MI,
         ++SrcIdx;
       }
     }
-    unsigned MaskIdx = SrcIdx + 1 + X86::AddrDisp;
-
-    assert(MI->getNumOperands() >= (SrcIdx + 1 + X86::AddrNumOperands) &&
-           "Unexpected number of operands!");
 
-    const MachineOperand &MaskOp = MI->getOperand(MaskIdx);
-    if (auto *C = X86::getConstantFromPool(*MI, MaskOp)) {
+    if (auto *C = X86::getConstantFromPool(*MI, SrcIdx + 1)) {
       unsigned Width = getRegisterWidth(MI->getDesc().operands()[0]);
       SmallVector<int, 64> Mask;
       DecodePSHUFBMask(C, Width, Mask);
@@ -1849,13 +1834,8 @@ static void addConstantComments(const MachineInstr *MI,
         ++SrcIdx;
       }
     }
-    unsigned MaskIdx = SrcIdx + 1 + X86::AddrDisp;
-
-    assert(MI->getNumOperands() >= (SrcIdx + 1 + X86::AddrNumOperands) &&
-           "Unexpected number of operands!");
 
-    const MachineOperand &MaskOp = MI->getOperand(MaskIdx);
-    if (auto *C = X86::getConstantFromPool(*MI, MaskOp)) {
+    if (auto *C = X86::getConstantFromPool(*MI, SrcIdx + 1)) {
       unsigned Width = getRegisterWidth(MI->getDesc().operands()[0]);
       SmallVector<int, 16> Mask;
       DecodeVPERMILPMask(C, ElSize, Width, Mask);
@@ -1883,8 +1863,7 @@ static void addConstantComments(const MachineInstr *MI,
     case X86::VPERMIL2PDrm: case X86::VPERMIL2PDYrm: ElSize = 64; break;
     }
 
-    const MachineOperand &MaskOp = MI->getOperand(3 + X86::AddrDisp);
-    if (auto *C = X86::getConstantFromPool(*MI, MaskOp)) {
+    if (auto *C = X86::getConstantFromPool(*MI, 3)) {
       unsigned Width = getRegisterWidth(MI->getDesc().operands()[0]);
       SmallVector<int, 16> Mask;
       DecodeVPERMIL2PMask(C, (unsigned)CtrlOp.getImm(), ElSize, Width, Mask);
@@ -1895,11 +1874,7 @@ static void addConstantComments(const MachineInstr *MI,
   }
 
   case X86::VPPERMrrm: {
-    assert(MI->getNumOperands() >= (3 + X86::AddrNumOperands) &&
-           "Unexpected number of operands!");
-
-    const MachineOperand &MaskOp = MI->getOperand(3 + X86::AddrDisp);
-    if (auto *C = X86::getConstantFromPool(*MI, MaskOp)) {
+    if (auto *C = X86::getConstantFromPool(*MI, 3)) {
       unsigned Width = getRegisterWidth(MI->getDesc().operands()[0]);
       SmallVector<int, 16> Mask;
       DecodeVPPERMMask(C, Width, Mask);
@@ -1910,10 +1885,7 @@ static void addConstantComments(const MachineInstr *MI,
   }
 
   case X86::MMX_MOVQ64rm: {
-    assert(MI->getNumOperands() == (1 + X86::AddrNumOperands) &&
-           "Unexpected number of operands!");
-    if (auto *C =
-            X86::getConstantFromPool(*MI, MI->getOperand(1 + X86::AddrDisp))) {
+    if (auto *C = X86::getConstantFromPool(*MI, 1)) {
       std::string Comment;
       raw_string_ostream CS(Comment);
       const MachineOperand &DstOp = MI->getOperand(0);