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

Amir Ayupov via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Thu Jul 18 18:02:12 PDT 2024


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

>From dd4d0de42048c063d5e5095a0c2594c7cc578df5 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 9 May 2024 19:35:26 -0700
Subject: [PATCH 1/5] Fix RISCVMCPlusBuilder

Created using spr 1.3.4
---
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index 74f2f0aae91e6..020e62463ee2f 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -177,13 +177,14 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
       MCInst &Instruction, InstructionIterator Begin, InstructionIterator End,
       const unsigned PtrSize, MCInst *&MemLocInstr, unsigned &BaseRegNum,
       unsigned &IndexRegNum, int64_t &DispValue, const MCExpr *&DispExpr,
-      MCInst *&PCRelBaseOut) const override {
+      MCInst *&PCRelBaseOut, MCInst *&FixedEntryLoadInst) const override {
     MemLocInstr = nullptr;
     BaseRegNum = 0;
     IndexRegNum = 0;
     DispValue = 0;
     DispExpr = nullptr;
     PCRelBaseOut = nullptr;
+    FixedEntryLoadInst = nullptr;
 
     // Check for the following long tail call sequence:
     // 1: auipc xi, %pcrel_hi(sym)

>From 62391bb5aa01f2b77d4315d1e72a9924eec9ecc0 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Fri, 5 Jul 2024 14:54:51 -0700
Subject: [PATCH 2/5] Drop deregisterJumpTable

Created using spr 1.3.4
---
 bolt/lib/Core/BinaryFunction.cpp | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 09a6ca1d68730..f587d5a2cadd4 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -899,17 +899,9 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
 
     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);
+    JumpTable *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);

>From 5336879ab68aedb1217e2c6c139d171f31e89e03 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Sat, 6 Jul 2024 22:26:14 -0700
Subject: [PATCH 3/5] Surgically drop spurious jump table

Created using spr 1.3.4
---
 bolt/include/bolt/Core/BinaryContext.h      |  5 +++++
 bolt/lib/Core/BinaryFunction.cpp            | 12 ++++++++++--
 bolt/test/X86/jump-table-fixed-ref-pic.test | 11 ++++-------
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 73932c4ca2fb3..c5e2c6cd02179 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -431,6 +431,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/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f587d5a2cadd4..2ecca32a5985c 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -899,9 +899,17 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
 
     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.
-    JumpTable *JT = BC.getJumpTableContainingAddress(ArrayStart);
+    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);
@@ -1158,10 +1166,10 @@ void BinaryFunction::handleIndirectBranch(MCInst &Instruction, uint64_t Size,
   }
   case IndirectBranchType::POSSIBLE_JUMP_TABLE:
   case IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE:
+  case IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH:
     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);
diff --git a/bolt/test/X86/jump-table-fixed-ref-pic.test b/bolt/test/X86/jump-table-fixed-ref-pic.test
index d43c9583f0d07..d215c565b31e5 100644
--- a/bolt/test/X86/jump-table-fixed-ref-pic.test
+++ b/bolt/test/X86/jump-table-fixed-ref-pic.test
@@ -7,10 +7,7 @@ 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: 			.Ltmp1 (2 instructions, align : 1)
-CHECK-NEXT: Secondary Entry Point: __ENTRY_main at 0x[[#TGT]]
+CHECK:      movslq ".rodata/1"+8(%rip), %rax
+CHECK-NEXT: leaq ".rodata/1"(%rip), %rdx
+CHECK-NEXT: addq %rdx, %rax
+CHECK-NEXT: jmpq *%rax # UNKNOWN CONTROL FLOW

>From e80b4ddfdb45f40a6d286b959f89f8fff818b5d7 Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Sun, 14 Jul 2024 14:50:22 -0700
Subject: [PATCH 4/5] Address comments

Created using spr 1.3.4
---
 bolt/include/bolt/Core/BinaryContext.h   |  6 +-
 bolt/lib/Core/BinaryContext.cpp          | 10 +++
 bolt/lib/Core/BinaryFunction.cpp         | 15 ++--
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp | 98 +++++++++++-------------
 4 files changed, 64 insertions(+), 65 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index c5e2c6cd02179..befdc9974c683 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -431,10 +431,8 @@ class BinaryContext {
     return nullptr;
   }
 
-  /// Deregister JumpTable registered at a given \p Address.
-  bool deregisterJumpTable(uint64_t Address) {
-    return JumpTables.erase(Address);
-  }
+  /// Deregister JumpTable registered at a given \p Address and delete it.
+  void deleteJumpTable(uint64_t Address);
 
   unsigned getDWARFEncodingSize(unsigned Encoding) {
     if (Encoding == dwarf::DW_EH_PE_omit)
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 0a1f1bb9e0d20..23b8eb266612e 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -2523,6 +2523,16 @@ BinaryFunction *BinaryContext::getBinaryFunctionAtAddress(uint64_t Address) {
   return nullptr;
 }
 
+/// Deregister JumpTable registered at a given \p Address and delete it.
+void BinaryContext::deleteJumpTable(uint64_t Address) {
+  JumpTable *JT = getJumpTableContainingAddress(Address);
+  assert(JT && "Must have a jump table at address");
+  for (BinaryFunction *Parent : JT->Parents)
+    Parent->JumpTables.erase(Address);
+  JumpTables.erase(Address);
+  delete JT;
+}
+
 DebugAddressRangesVector BinaryContext::translateModuleAddressRanges(
     const DWARFAddressRangesVector &InputRanges) const {
   DebugAddressRangesVector OutputRanges;
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 2ecca32a5985c..c0d570d6d909c 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -883,8 +883,11 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
   if (FixedEntryLoadInstr) {
     assert(BranchType == IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH &&
            "Invalid IndirectBranch type");
-    const MCExpr *FixedEntryDispExpr =
-        BC.MIB->getMemOperandDisp(*FixedEntryLoadInstr)->getExpr();
+    MCInst::iterator FixedEntryDispOperand =
+        BC.MIB->getMemOperandDisp(*FixedEntryLoadInstr);
+    assert(FixedEntryDispOperand != FixedEntryLoadInstr->end() &&
+           "Invalid memory instruction");
+    const MCExpr *FixedEntryDispExpr = FixedEntryDispOperand->getExpr();
     const uint64_t EntryAddress = getExprValue(FixedEntryDispExpr);
     uint64_t EntrySize = BC.getJumpTableEntrySize(JumpTable::JTT_PIC);
     ErrorOr<int64_t> Value = BC.getSignedValueAtAddress(EntryAddress, EntrySize);
@@ -901,15 +904,11 @@ BinaryFunction::processIndirectBranch(MCInst &Instruction, unsigned Size,
 
     // 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;
+    BC.deleteJumpTable(EntryAddress);
 
     // Replace FixedEntryDispExpr used in target address calculation with outer
     // jump table reference.
-    JT = BC.getJumpTableContainingAddress(ArrayStart);
+    JumpTable *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);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 527ef5932c550..604ed7f3b00b7 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1923,14 +1923,14 @@ class X86MCPlusBuilder : public MCPlusBuilder {
              MO.SegRegNum == X86::NoRegister;
     };
     LLVM_DEBUG(dbgs() << "Checking for PIC jump table\n");
-    MCInst *MemLocInstr = nullptr;
-    MCInst *MovInstr = nullptr;
-    bool IsFixedBranch = false;
-    // Allow renaming R1/R2 once.
-    bool SwappedRegs = false;
+    MCInst *FirstInstr = nullptr;
+    MCInst *SecondInstr = nullptr;
+    enum {
+      NOMATCH = 0,
+      MATCH_JUMP_TABLE,
+      MATCH_FIXED_BRANCH,
+    } MatchingState = NOMATCH;
     while (++II != IE) {
-      if (MovInstr && MemLocInstr)
-        break;
       MCInst &Instr = *II;
       const MCInstrDesc &InstrDesc = Info->get(Instr.getOpcode());
       if (!InstrDesc.hasDefOfPhysReg(Instr, R1, *RegInfo) &&
@@ -1938,75 +1938,67 @@ class X86MCPlusBuilder : public MCPlusBuilder {
         // Ignore instructions that don't affect R1, R2 registers.
         continue;
       }
-      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 (!SwappedRegs && MovDestReg != R2) {
+      const bool IsMOVSXInstr = isMOVSX64rm32(Instr);
+      const bool IsLEAInstr = isLEA64r(Instr);
+      if (MatchingState == NOMATCH) {
+        if (IsMOVSXInstr)
+          MatchingState = MATCH_JUMP_TABLE;
+        else if (IsLEAInstr)
+          MatchingState = MATCH_FIXED_BRANCH;
+        else
+          break;
+
+        // Check if the first instruction is setting %r1 or %r2. In canonical
+        // form lea sets %r1 and mov sets %r2. If it's the opposite - rename so
+        // we have to only check a single form.
+        unsigned DestReg = Instr.getOperand(0).getReg();
+        MCPhysReg &ExpectReg = MatchingState == MATCH_JUMP_TABLE ? R2 : R1;
+        if (DestReg != ExpectReg)
           std::swap(R1, R2);
-          SwappedRegs = true;
-        }
-        if (MovDestReg != R2) {
-          LLVM_DEBUG(dbgs() << "MOV instruction expected to set %r2\n");
+        if (DestReg != ExpectReg)
           break;
-        }
 
-        // Verify operands for MOV.
+        // Verify operands
         std::optional<X86MemOperand> MO = evaluateX86MemoryOperand(Instr);
         if (!MO)
           break;
-        if (isRIPRel(*MO))
-          IsFixedBranch = true;
-        else if (isIndexed(*MO, R1))
-          ; // POSSIBLE_PIC_JUMP_TABLE
+        if ((MatchingState == MATCH_JUMP_TABLE && isIndexed(*MO, R1)) ||
+            (MatchingState == MATCH_FIXED_BRANCH && isRIPRel(*MO)))
+          FirstInstr = &Instr;
         else
           break;
-        MovInstr = &Instr;
-        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");
+      } else {
+        unsigned ExpectReg = MatchingState == MATCH_JUMP_TABLE ? R1 : R2;
+        if (!InstrDesc.hasDefOfPhysReg(Instr, ExpectReg, *RegInfo))
+          continue;
+        if ((MatchingState == MATCH_JUMP_TABLE && !IsLEAInstr) ||
+            (MatchingState == MATCH_FIXED_BRANCH && !IsMOVSXInstr))
+          break;
+        if (Instr.getOperand(0).getReg() != ExpectReg)
           break;
-        }
 
-        // Verify operands for LEA.
+        // Verify operands.
         std::optional<X86MemOperand> MO = evaluateX86MemoryOperand(Instr);
         if (!MO)
           break;
         if (!isRIPRel(*MO))
           break;
-        MemLocInstr = &Instr;
-        continue;
+        SecondInstr = &Instr;
+        break;
       }
     }
 
-    if (!MemLocInstr)
+    if (!SecondInstr)
       return std::make_tuple(IndirectBranchType::UNKNOWN, nullptr, nullptr);
 
-    LLVM_DEBUG(dbgs() << "checking potential fixed indirect branch\n");
-    if (IsFixedBranch)
+    if (MatchingState == MATCH_FIXED_BRANCH) {
+      LLVM_DEBUG(dbgs() << "checking potential fixed indirect branch\n");
       return std::make_tuple(IndirectBranchType::POSSIBLE_PIC_FIXED_BRANCH,
-                             MemLocInstr, MovInstr);
-
+                             FirstInstr, SecondInstr);
+    }
     LLVM_DEBUG(dbgs() << "checking potential PIC jump table\n");
     return std::make_tuple(IndirectBranchType::POSSIBLE_PIC_JUMP_TABLE,
-                          MemLocInstr, nullptr);
+                           SecondInstr, nullptr);
   }
 
   IndirectBranchType

>From 677cb3248aaae5c441cceff074e3569f84556cdb Mon Sep 17 00:00:00 2001
From: Amir Ayupov <aaupov at fb.com>
Date: Thu, 18 Jul 2024 18:02:00 -0700
Subject: [PATCH 5/5] Address comments

Created using spr 1.3.4
---
 bolt/lib/Core/BinaryContext.cpp          | 4 ++--
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp | 5 ++++-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index 23b8eb266612e..035f68e39751b 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -2525,8 +2525,8 @@ BinaryFunction *BinaryContext::getBinaryFunctionAtAddress(uint64_t Address) {
 
 /// Deregister JumpTable registered at a given \p Address and delete it.
 void BinaryContext::deleteJumpTable(uint64_t Address) {
-  JumpTable *JT = getJumpTableContainingAddress(Address);
-  assert(JT && "Must have a jump table at address");
+  assert(JumpTables.count(Address) && "Must have a jump table at address");
+  JumpTable *JT = JumpTables.at(Address);
   for (BinaryFunction *Parent : JT->Parents)
     Parent->JumpTables.erase(Address);
   JumpTables.erase(Address);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 604ed7f3b00b7..9f6c552977b34 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -1866,6 +1866,9 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     return true;
   }
 
+  /// Analyzes PIC-style jump table code template and return identified
+  /// IndirectBranchType, MemLocInstr (all cases) and FixedEntryLoadInstr
+  /// (POSSIBLE_PIC_FIXED_BRANCH case).
   template <typename Itr>
   std::tuple<IndirectBranchType, MCInst *, MCInst *>
   analyzePICJumpTable(Itr II, Itr IE, MCPhysReg R1, MCPhysReg R2) const {
@@ -1878,7 +1881,7 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     //
     // or a fixed indirect jump template:
     //
-    //    movslq En(%rip), {%r2|%r1}
+    //    movslq En(%rip), {%r2|%r1}              <- FixedEntryLoadInstr
     //    lea PIC_JUMP_TABLE(%rip), {%r1|%r2}     <- MemLocInstr
     //    add %r2, %r1
     //    jmp *%r1



More information about the llvm-branch-commits mailing list