[llvm-branch-commits] [BOLT] Support POSSIBLE_PIC_FIXED_BRANCH (PR #91667)

via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Thu May 9 14:56:52 PDT 2024


llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-bolt

Author: Amir Ayupov (aaupov)

<details>
<summary>Changes</summary>

Detect and support fixed PIC indirect jumps of the following form:
```
movslq  En(%rip), %r1
leaq  PIC_JUMP_TABLE(%rip), %r2
addq  %r2, %r1
jmpq  *%r1
```

with PIC_JUMP_TABLE that looks like following:

```
  JT:  ----------
   E1:| L1 - JT  |
      |----------|
   E2:| L2 - JT  |
      |----------|
      |          |
         ......
   En:| Ln - JT  |
       ----------
```

The code could be produced by compilers, see
https://github.com/llvm/llvm-project/issues/91648.

Test Plan: updated jump-table-fixed-ref-pic.test


---
Full diff: https://github.com/llvm/llvm-project/pull/91667.diff


8 Files Affected:

- (modified) bolt/include/bolt/Core/BinaryContext.h (+5) 
- (modified) bolt/include/bolt/Core/MCPlusBuilder.h (+4-1) 
- (modified) bolt/lib/Core/BinaryFunction.cpp (+45-2) 
- (modified) bolt/lib/Passes/IndirectCallPromotion.cpp (+2-1) 
- (modified) bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp (+9-5) 
- (modified) bolt/lib/Target/X86/X86MCPlusBuilder.cpp (+60-28) 
- (modified) bolt/test/X86/Inputs/jump-table-fixed-ref-pic.s (+1-1) 
- (modified) bolt/test/X86/jump-table-fixed-ref-pic.test (+11-4) 


``````````diff
diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 4a59a581dfedb..f8bf29c674b54 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -430,6 +430,11 @@ class BinaryContext {
     return nullptr;
   }
 
+  /// Deregister JumpTable registered at a given \p Address.
+  bool deregisterJumpTable(uint64_t Address) {
+    return JumpTables.erase(Address);
+  }
+
   unsigned getDWARFEncodingSize(unsigned Encoding) {
     if (Encoding == dwarf::DW_EH_PE_omit)
       return 0;
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index f7614cf9ac977..42ec006fba9f1 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -58,6 +58,8 @@ enum class IndirectBranchType : char {
   POSSIBLE_PIC_JUMP_TABLE, /// Possibly a jump table for PIC.
   POSSIBLE_GOTO,           /// Possibly a gcc's computed goto.
   POSSIBLE_FIXED_BRANCH,   /// Possibly an indirect branch to a fixed location.
+  POSSIBLE_PIC_FIXED_BRANCH, /// Possibly an indirect jump to a fixed entry in a
+                             /// PIC jump table.
 };
 
 class MCPlusBuilder {
@@ -1472,7 +1474,8 @@ class MCPlusBuilder {
                         InstructionIterator End, const unsigned PtrSize,
                         MCInst *&MemLocInstr, unsigned &BaseRegNum,
                         unsigned &IndexRegNum, int64_t &DispValue,
-                        const MCExpr *&DispExpr, MCInst *&PCRelBaseOut) const {
+                        const MCExpr *&DispExpr, MCInst *&PCRelBaseOut,
+                        MCInst *&FixedEntryLoadInst) const {
     llvm_unreachable("not implemented");
     return IndirectBranchType::UNKNOWN;
   }
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f74ecea8ac0a1..799065a6f194e 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -779,6 +779,9 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
   // setting the value of the register used by the branch.
   MCInst *MemLocInstr;
 
+  // The instruction loading the fixed PIC jump table entry value.
+  MCInst *FixedEntryLoadInstr;
+
   // Address of the table referenced by MemLocInstr. Could be either an
   // array of function pointers, or a jump table.
   uint64_t ArrayStart = 0;
@@ -810,7 +813,7 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
 
   IndirectBranchType BranchType = BC.MIB->analyzeIndirectBranch(
       Instruction, Begin, Instructions.end(), PtrSize, MemLocInstr, BaseRegNum,
-      IndexRegNum, DispValue, DispExpr, PCRelBaseInstr);
+      IndexRegNum, DispValue, DispExpr, PCRelBaseInstr, FixedEntryLoadInstr);
 
   if (BranchType == IndirectBranchType::UNKNOWN && !MemLocInstr)
     return BranchType;
@@ -876,6 +879,43 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
   if (BaseRegNum == BC.MRI->getProgramCounter())
     ArrayStart += getAddress() + Offset + Size;
 
+  if (FixedEntryLoadInstr) {
+    assert(BranchType == IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH &&
+           "Invalid IndirectBranch type");
+    const MCExpr *FixedEntryDispExpr =
+        BC.MIB->getMemOperandDisp(*FixedEntryLoadInstr)->getExpr();
+    const uint64_t EntryAddress = getExprValue(FixedEntryDispExpr);
+    uint64_t EntrySize = BC.getJumpTableEntrySize(JumpTable::JTT_PIC);
+    ErrorOr<int64_t> Value = BC.getSignedValueAtAddress(EntryAddress, EntrySize);
+    if (!Value)
+      return IndirectBranchType::UNKNOWN;
+
+    BC.outs() << "BOLT-INFO: fixed PIC indirect branch detected in " << *this
+              << " at 0x" << Twine::utohexstr(getAddress() + Offset)
+              << " referencing data at 0x" << Twine::utohexstr(EntryAddress)
+              << " the destination value is 0x"
+              << Twine::utohexstr(ArrayStart + *Value) << '\n';
+
+    TargetAddress = ArrayStart + *Value;
+
+    // Remove spurious JumpTable at EntryAddress caused by PIC reference from
+    // the load instruction.
+    JumpTable *JT = BC.getJumpTableContainingAddress(EntryAddress);
+    assert(JT && "Must have a jump table at fixed entry address");
+    BC.deregisterJumpTable(EntryAddress);
+    JumpTables.erase(EntryAddress);
+    delete JT;
+
+    // Replace FixedEntryDispExpr used in target address calculation with outer
+    // jump table reference.
+    JT = BC.getJumpTableContainingAddress(ArrayStart);
+    assert(JT && "Must have a containing jump table for PIC fixed branch");
+    BC.MIB->replaceMemOperandDisp(*FixedEntryLoadInstr, JT->getFirstLabel(),
+                                  EntryAddress - ArrayStart, &*BC.Ctx);
+
+    return BranchType;
+  }
+
   LLVM_DEBUG(dbgs() << "BOLT-DEBUG: addressed memory is 0x"
                     << Twine::utohexstr(ArrayStart) << '\n');
 
@@ -1128,6 +1168,7 @@ void BinaryFunction::handleIndirectBranch(MCInst &Instruction, uint64_t Size,
     if (opts::JumpTables == JTS_NONE)
       IsSimple = false;
     break;
+  case IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH:
   case IndirectBranchType::POSSIBLE_FIXED_BRANCH: {
     if (containsAddress(IndirectTarget)) {
       const MCSymbol *TargetSymbol = getOrCreateLocalLabel(IndirectTarget);
@@ -1894,9 +1935,11 @@ bool BinaryFunction::postProcessIndirectBranches(
         int64_t DispValue;
         const MCExpr *DispExpr;
         MCInst *PCRelBaseInstr;
+        MCInst *FixedEntryLoadInstr;
         IndirectBranchType Type = BC.MIB->analyzeIndirectBranch(
             Instr, BB.begin(), II, PtrSize, MemLocInstr, BaseRegNum,
-            IndexRegNum, DispValue, DispExpr, PCRelBaseInstr);
+            IndexRegNum, DispValue, DispExpr, PCRelBaseInstr,
+            FixedEntryLoadInstr);
         if (Type != IndirectBranchType::UNKNOWN || MemLocInstr != nullptr)
           continue;
 
diff --git a/bolt/lib/Passes/IndirectCallPromotion.cpp b/bolt/lib/Passes/IndirectCallPromotion.cpp
index 55eede641fd2f..9483f7c86e366 100644
--- a/bolt/lib/Passes/IndirectCallPromotion.cpp
+++ b/bolt/lib/Passes/IndirectCallPromotion.cpp
@@ -386,13 +386,14 @@ IndirectCallPromotion::maybeGetHotJumpTableTargets(BinaryBasicBlock &BB,
   JumpTableInfoType HotTargets;
   MCInst *MemLocInstr;
   MCInst *PCRelBaseOut;
+  MCInst *FixedEntryLoadInstr;
   unsigned BaseReg, IndexReg;
   int64_t DispValue;
   const MCExpr *DispExpr;
   MutableArrayRef<MCInst> Insts(&BB.front(), &CallInst);
   const IndirectBranchType Type = BC.MIB->analyzeIndirectBranch(
       CallInst, Insts.begin(), Insts.end(), BC.AsmInfo->getCodePointerSize(),
-      MemLocInstr, BaseReg, IndexReg, DispValue, DispExpr, PCRelBaseOut);
+      MemLocInstr, BaseReg, IndexReg, DispValue, DispExpr, FixedEntryLoadInstr);
 
   assert(MemLocInstr && "There should always be a load for jump tables");
   if (!MemLocInstr)
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 0ae9d3668b93b..8fa004b817579 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -832,16 +832,20 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return Uses;
   }
 
-  IndirectBranchType analyzeIndirectBranch(
-      MCInst &Instruction, InstructionIterator Begin, InstructionIterator End,
-      const unsigned PtrSize, MCInst *&MemLocInstrOut, unsigned &BaseRegNumOut,
-      unsigned &IndexRegNumOut, int64_t &DispValueOut,
-      const MCExpr *&DispExprOut, MCInst *&PCRelBaseOut) const override {
+  IndirectBranchType
+  analyzeIndirectBranch(MCInst &Instruction, InstructionIterator Begin,
+                        InstructionIterator End, const unsigned PtrSize,
+                        MCInst *&MemLocInstrOut, unsigned &BaseRegNumOut,
+                        unsigned &IndexRegNumOut, int64_t &DispValueOut,
+                        const MCExpr *&DispExprOut, MCInst *&PCRelBaseOut,
+                        const MCExpr *&JTBaseDispExprOut, MCInst *&PCRelBaseOut,
+                        MCInst *&FixedEntryLoadInstr) const override {
     MemLocInstrOut = nullptr;
     BaseRegNumOut = AArch64::NoRegister;
     IndexRegNumOut = AArch64::NoRegister;
     DispValueOut = 0;
     DispExprOut = nullptr;
+    FixedEntryLoadInstr = nullptr;
 
     // An instruction referencing memory used by jump instruction (directly or
     // via register). This location could be an array of function pointers
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 86e7d4dfaed8d..f7c173aa23b35 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1897,7 +1897,7 @@ class X86MCPlusBuilder : public MCPlusBuilder {
   }
 
   template <typename Itr>
-  std::pair<IndirectBranchType, MCInst *>
+  std::tuple<IndirectBranchType, MCInst *, MCInst *>
   analyzePICJumpTable(Itr II, Itr IE, MCPhysReg R1, MCPhysReg R2) const {
     // Analyze PIC-style jump table code template:
     //
@@ -1906,6 +1906,13 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     //    add %r2, %r1
     //    jmp *%r1
     //
+    // or a fixed indirect jump template:
+    //
+    //    movslq En(%rip), {%r2|%r1}
+    //    lea PIC_JUMP_TABLE(%rip), {%r1|%r2}     <- MemLocInstr
+    //    add %r2, %r1
+    //    jmp *%r1
+    //
     // (with any irrelevant instructions in-between)
     //
     // When we call this helper we've already determined %r1 and %r2, and
@@ -1947,8 +1954,13 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     };
     LLVM_DEBUG(dbgs() << "Checking for PIC jump table\n");
     MCInst *MemLocInstr = nullptr;
-    const MCInst *MovInstr = nullptr;
+    MCInst *MovInstr = nullptr;
+    bool IsFixedBranch = false;
+    // Allow renaming R1/R2 once.
+    bool SwappedRegs = false;
     while (++II != IE) {
+      if (MovInstr && MemLocInstr)
+        break;
       MCInst &Instr = *II;
       const MCInstrDesc &InstrDesc = Info->get(Instr.getOpcode());
       if (!InstrDesc.hasDefOfPhysReg(Instr, R1, *RegInfo) &&
@@ -1956,19 +1968,18 @@ class X86MCPlusBuilder : public MCPlusBuilder {
         // Ignore instructions that don't affect R1, R2 registers.
         continue;
       }
-      if (!MovInstr) {
-        // Expect to see MOV instruction.
-        if (!isMOVSX64rm32(Instr)) {
-          LLVM_DEBUG(dbgs() << "MOV instruction expected.\n");
-          break;
-        }
-
+      if (isMOVSX64rm32(Instr)) {
+        // Potential redefinition of MovInstr - bail.
+        if (MovInstr)
+          return std::make_tuple(IndirectBranchType::UNKNOWN, nullptr, nullptr);
         // Check if it's setting %r1 or %r2. In canonical form it sets %r2.
         // If it sets %r1 - rename the registers so we have to only check
         // a single form.
         unsigned MovDestReg = Instr.getOperand(0).getReg();
-        if (MovDestReg != R2)
+        if (!SwappedRegs && MovDestReg != R2) {
           std::swap(R1, R2);
+          SwappedRegs = true;
+        }
         if (MovDestReg != R2) {
           LLVM_DEBUG(dbgs() << "MOV instruction expected to set %r2\n");
           break;
@@ -1978,16 +1989,26 @@ class X86MCPlusBuilder : public MCPlusBuilder {
         std::optional<X86MemOperand> MO = evaluateX86MemoryOperand(Instr);
         if (!MO)
           break;
-        if (!isIndexed(*MO, R1))
-          // POSSIBLE_PIC_JUMP_TABLE
+        if (isRIPRel(*MO))
+          IsFixedBranch = true;
+        else if (isIndexed(*MO, R1))
+          ; // POSSIBLE_PIC_JUMP_TABLE
+        else
           break;
         MovInstr = &Instr;
-      } else {
-        if (!InstrDesc.hasDefOfPhysReg(Instr, R1, *RegInfo))
-          continue;
-        if (!isLEA64r(Instr)) {
-          LLVM_DEBUG(dbgs() << "LEA instruction expected\n");
-          break;
+        continue;
+      }
+      if (isLEA64r(Instr)) {
+        // Potential redefinition of MemLocInstr - bail.
+        if (MemLocInstr)
+          return std::make_tuple(IndirectBranchType::UNKNOWN, nullptr, nullptr);
+        // Check if it's setting %r1 or %r2. In canonical form it sets %r1.
+        // If it sets %r2 - rename the registers so we have to only check
+        // a single form.
+        unsigned LeaDestReg = Instr.getOperand(0).getReg();
+        if (!SwappedRegs && LeaDestReg != R1) {
+          std::swap(R1, R2);
+          SwappedRegs = true;
         }
         if (Instr.getOperand(0).getReg() != R1) {
           LLVM_DEBUG(dbgs() << "LEA instruction expected to set %r1\n");
@@ -2001,23 +2022,30 @@ class X86MCPlusBuilder : public MCPlusBuilder {
         if (!isRIPRel(*MO))
           break;
         MemLocInstr = &Instr;
-        break;
+        continue;
       }
     }
 
     if (!MemLocInstr)
-      return std::make_pair(IndirectBranchType::UNKNOWN, nullptr);
+      return std::make_tuple(IndirectBranchType::UNKNOWN, nullptr, nullptr);
+
+    LLVM_DEBUG(dbgs() << "checking potential fixed indirect branch\n");
+    if (IsFixedBranch)
+      return std::make_tuple(IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH,
+                             MemLocInstr, MovInstr);
 
     LLVM_DEBUG(dbgs() << "checking potential PIC jump table\n");
-    return std::make_pair(IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE,
-                          MemLocInstr);
+    return std::make_tuple(IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE,
+                          MemLocInstr, nullptr);
   }
 
-  IndirectBranchType analyzeIndirectBranch(
-      MCInst &Instruction, InstructionIterator Begin, InstructionIterator End,
-      const unsigned PtrSize, MCInst *&MemLocInstrOut, unsigned &BaseRegNumOut,
-      unsigned &IndexRegNumOut, int64_t &DispValueOut,
-      const MCExpr *&DispExprOut, MCInst *&PCRelBaseOut) const override {
+  IndirectBranchType
+  analyzeIndirectBranch(MCInst &Instruction, InstructionIterator Begin,
+                        InstructionIterator End, const unsigned PtrSize,
+                        MCInst *&MemLocInstrOut, unsigned &BaseRegNumOut,
+                        unsigned &IndexRegNumOut, int64_t &DispValueOut,
+                        const MCExpr *&DispExprOut, MCInst *&PCRelBaseOut,
+                        MCInst *&FixedEntryLoadInst) const override {
     // Try to find a (base) memory location from where the address for
     // the indirect branch is loaded. For X86-64 the memory will be specified
     // in the following format:
@@ -2044,6 +2072,7 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     IndexRegNumOut = X86::NoRegister;
     DispValueOut = 0;
     DispExprOut = nullptr;
+    FixedEntryLoadInst = nullptr;
 
     std::reverse_iterator<InstructionIterator> II(End);
     std::reverse_iterator<InstructionIterator> IE(Begin);
@@ -2076,7 +2105,8 @@ class X86MCPlusBuilder : public MCPlusBuilder {
           unsigned R2 = PrevInstr.getOperand(2).getReg();
           if (R1 == R2)
             return IndirectBranchType::UNKNOWN;
-          std::tie(Type, MemLocInstr) = analyzePICJumpTable(PrevII, IE, R1, R2);
+          std::tie(Type, MemLocInstr, FixedEntryLoadInst) =
+              analyzePICJumpTable(PrevII, IE, R1, R2);
           break;
         }
         return IndirectBranchType::UNKNOWN;
@@ -2120,6 +2150,8 @@ class X86MCPlusBuilder : public MCPlusBuilder {
       if (MO->ScaleImm != 1 || MO->BaseRegNum != RIPRegister)
         return IndirectBranchType::UNKNOWN;
       break;
+    case IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH:
+      break;
     default:
       if (MO->ScaleImm != PtrSize)
         return IndirectBranchType::UNKNOWN;
diff --git a/bolt/test/X86/Inputs/jump-table-fixed-ref-pic.s b/bolt/test/X86/Inputs/jump-table-fixed-ref-pic.s
index 66629a4880e64..6407964593e2d 100644
--- a/bolt/test/X86/Inputs/jump-table-fixed-ref-pic.s
+++ b/bolt/test/X86/Inputs/jump-table-fixed-ref-pic.s
@@ -6,7 +6,7 @@ main:
   jae .L4
   cmpq $0x1, %rdi
   jne .L4
-  mov .Ljt_pic+8(%rip), %rax
+  movslq .Ljt_pic+8(%rip), %rax
   lea .Ljt_pic(%rip), %rdx
   add %rdx, %rax
   jmpq *%rax
diff --git a/bolt/test/X86/jump-table-fixed-ref-pic.test b/bolt/test/X86/jump-table-fixed-ref-pic.test
index 4195b97aac501..d2e90eb1d4099 100644
--- a/bolt/test/X86/jump-table-fixed-ref-pic.test
+++ b/bolt/test/X86/jump-table-fixed-ref-pic.test
@@ -1,9 +1,16 @@
 # Verify that BOLT detects fixed destination of indirect jump for PIC
 # case.
 
-XFAIL: *
-
 RUN: %clang %cflags -no-pie %S/Inputs/jump-table-fixed-ref-pic.s -Wl,-q -o %t
-RUN: llvm-bolt %t --relocs -o %t.null 2>&1 | FileCheck %s
+RUN: llvm-bolt %t --relocs -o %t.null -print-cfg 2>&1 | FileCheck %s
+
+CHECK: BOLT-INFO: fixed PIC indirect branch detected in main {{.*}} the destination value is 0x[[#TGT:]]
+CHECK: Binary Function "main" after building cfg
+
+CHECK: 			movslq  ".rodata/1"+8(%rip), %rax
+CHECK-NEXT: leaq    ".rodata/1"(%rip), %rdx
+CHECK-NEXT: addq    %rdx, %rax
+CHECK-NEXT: jmp     .Ltmp1
 
-CHECK: BOLT-INFO: fixed indirect branch detected in main
+CHECK: 			.Ltmp1 (2 instructions, align : 1)
+CHECK-NEXT: Secondary Entry Point: __ENTRY_main at 0x[[#TGT]]

``````````

</details>


https://github.com/llvm/llvm-project/pull/91667


More information about the llvm-branch-commits mailing list