[X86] Add MI-layer routine for getting the index of the first address operand, NFC (#78019)

Add the MI-layer routine X86::getFirstAddrOperandIdx(), which returns
the index of the first address operand of a MachineInstr (or -1 if there
is none).

X86II::getMemoryOperandNo(), the existing MC-layer routine used to
obtain the index of the first address operand in a 5-operand X86 memory
reference, is incomplete: it does not handle pseudo-instructions like
TCRETURNmi, resulting in security holes in the mitigation passes that
use it (e.g., x86-slh and x86-lvi-load).

X86::getFirstAddrOperandIdx() handles both pseudo and real instructions
and is thus more suitable for most use cases than
X86II::getMemoryOperandNo(), especially in mitigation passes like
x86-slh and x86-lvi-load. For this reason, this patch replaces all uses
of X86II::getMemoryOperandNo() with X86::getFirstAddrOperandIdx() in the
aforementioned mitigation passes.

llvm/lib/Target/X86/X86InstrInfo.cpp

@@ -3463,6 +3463,56 @@ bool X86::isX87Instruction(MachineInstr &MI) {
   return false;
 }
 
+int X86::getFirstAddrOperandIdx(const MachineInstr &MI) {
+  const auto isMemOp = [](const MCOperandInfo &OpInfo) -> bool {
+    return OpInfo.OperandType == MCOI::OPERAND_MEMORY;
+  };
+
+  const MCInstrDesc &Desc = MI.getDesc();
+
+  // Directly invoke the MC-layer routine for real (i.e., non-pseudo)
+  // instructions (fast case).
+  if (!X86II::isPseudo(Desc.TSFlags)) {
+    int MemRefIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
+    if (MemRefIdx >= 0)
+      return MemRefIdx + X86II::getOperandBias(Desc);
+#ifdef EXPENSIVE_CHECKS
+    assert(none_of(Desc.operands(), isMemOp) &&
+           "Got false negative from X86II::getMemoryOperandNo()!");
+#endif
+    return -1;
+  }
+
+  // Otherwise, handle pseudo instructions by examining the type of their
+  // operands (slow case). An instruction cannot have a memory reference if it
+  // has fewer than AddrNumOperands (= 5) explicit operands.
+  if (Desc.getNumOperands() < X86::AddrNumOperands) {
+#ifdef EXPENSIVE_CHECKS
+    assert(none_of(Desc.operands(), isMemOp) &&
+           "Expected no operands to have OPERAND_MEMORY type!");
+#endif
+    return -1;
+  }
+
+  // The first operand with type OPERAND_MEMORY indicates the start of a memory
+  // reference. We expect the following AddrNumOperand-1 operands to also have
+  // OPERAND_MEMORY type.
+  for (unsigned i = 0; i <= Desc.getNumOperands() - X86::AddrNumOperands; ++i) {
+    if (Desc.operands()[i].OperandType == MCOI::OPERAND_MEMORY) {
+#ifdef EXPENSIVE_CHECKS
+      assert(std::all_of(Desc.operands().begin() + i,
+                         Desc.operands().begin() + i + X86::AddrNumOperands,
+                         isMemOp) &&
+             "Expected all five operands in the memory reference to have "
+             "OPERAND_MEMORY type!");
+#endif
+      return i;
+    }
+  }
+
+  return -1;
+}
+
 bool X86InstrInfo::isUnconditionalTailCall(const MachineInstr &MI) const {
   switch (MI.getOpcode()) {
   case X86::TCRETURNdi:

llvm/lib/Target/X86/X86InstrInfo.h

@@ -79,6 +79,13 @@ unsigned getSwappedVCMPImm(unsigned Imm);
 
 /// Check if the instruction is X87 instruction.
 bool isX87Instruction(MachineInstr &MI);
+
+/// Return the index of the instruction's first address operand, if it has a
+/// memory reference, or -1 if it has none. Unlike X86II::getMemoryOperandNo(),
+/// this also works for both pseudo instructions (e.g., TCRETURNmi) as well as
+/// real instructions (e.g., JMP64m).
+int getFirstAddrOperandIdx(const MachineInstr &MI);
+
 } // namespace X86
 
 /// isGlobalStubReference - Return true if the specified TargetFlag operand is

llvm/lib/Target/X86/X86LoadValueInjectionLoadHardening.cpp

@@ -770,16 +770,13 @@ bool X86LoadValueInjectionLoadHardeningPass::instrUsesRegToAccessMemory(
       MI.getOpcode() == X86::SFENCE || MI.getOpcode() == X86::LFENCE)
     return false;
 
-  // FIXME: This does not handle pseudo loading instruction like TCRETURN*
+  const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI);
-  const MCInstrDesc &Desc = MI.getDesc();
-  int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
   if (MemRefBeginIdx < 0) {
     LLVM_DEBUG(dbgs() << "Warning: unable to obtain memory operand for loading "
                          "instruction:\n";
                MI.print(dbgs()); dbgs() << '\n';);
     return false;
   }
-  MemRefBeginIdx += X86II::getOperandBias(Desc);
 
   const MachineOperand &BaseMO =
       MI.getOperand(MemRefBeginIdx + X86::AddrBaseReg);

llvm/lib/Target/X86/X86SpeculativeLoadHardening.cpp

@@ -1317,12 +1317,7 @@ void X86SpeculativeLoadHardeningPass::tracePredStateThroughBlocksAndHarden(
           continue;
 
         // Extract the memory operand information about this instruction.
-        // FIXME: This doesn't handle loading pseudo instructions which we often
+        const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI);
-        // could handle with similarly generic logic. We probably need to add an
-        // MI-layer routine similar to the MC-layer one we use here which maps
-        // pseudos much like this maps real instructions.
-        const MCInstrDesc &Desc = MI.getDesc();
-        int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
         if (MemRefBeginIdx < 0) {
           LLVM_DEBUG(dbgs()
                          << "WARNING: unable to harden loading instruction: ";
@@ -1330,8 +1325,6 @@ void X86SpeculativeLoadHardeningPass::tracePredStateThroughBlocksAndHarden(
           continue;
         }
 
-        MemRefBeginIdx += X86II::getOperandBias(Desc);
-
         MachineOperand &BaseMO =
             MI.getOperand(MemRefBeginIdx + X86::AddrBaseReg);
         MachineOperand &IndexMO =
@@ -1400,12 +1393,9 @@ void X86SpeculativeLoadHardeningPass::tracePredStateThroughBlocksAndHarden(
 
         // Check if this is a load whose address needs to be hardened.
         if (HardenLoadAddr.erase(&MI)) {
-          const MCInstrDesc &Desc = MI.getDesc();
+          const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(MI);
-          int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
           assert(MemRefBeginIdx >= 0 && "Cannot have an invalid index here!");
 
-          MemRefBeginIdx += X86II::getOperandBias(Desc);
-
           MachineOperand &BaseMO =
               MI.getOperand(MemRefBeginIdx + X86::AddrBaseReg);
           MachineOperand &IndexMO =
@@ -1802,11 +1792,9 @@ MachineInstr *X86SpeculativeLoadHardeningPass::sinkPostLoadHardenedInst(
 
         // Otherwise, this is a load and the load component can't be data
         // invariant so check how this register is being used.
-        const MCInstrDesc &Desc = UseMI.getDesc();
+        const int MemRefBeginIdx = X86::getFirstAddrOperandIdx(UseMI);
-        int MemRefBeginIdx = X86II::getMemoryOperandNo(Desc.TSFlags);
         assert(MemRefBeginIdx >= 0 &&
                "Should always have mem references here!");
-        MemRefBeginIdx += X86II::getOperandBias(Desc);
 
         MachineOperand &BaseMO =
             UseMI.getOperand(MemRefBeginIdx + X86::AddrBaseReg);