jumptable

Author: Ryan-hub-bit

llvm/lib/Target/X86/CMakeLists.txt

@@ -23,7 +23,7 @@ tablegen(LLVM X86GenFoldTables.inc -gen-x86-fold-tables -asmwriternum=1)
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
-X86MatchJumptablePass.cpp
+  X86MatchJumptablePass.cpp
   X86ArgumentStackSlotRebase.cpp
   X86AsmPrinter.cpp
   X86AvoidTrailingCall.cpp

llvm/lib/Target/X86/X86MatchJumptablePass.cpp

@@ -6,7 +6,7 @@
 #include "llvm/Support/raw_ostream.h"
 #include "X86MatchJumptablePass.h"
 
-#define DEBUG_TYPE "x86-my-pass"
+#define DEBUG_TYPE "match-jump-table"
 
 using namespace llvm;
 
@@ -18,23 +18,97 @@ namespace {
     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
 
     bool runOnMachineFunction(MachineFunction &MF) override {
-      LLVM_DEBUG(dbgs() << "Running X86MyBackendPass on function: "
-                        << MF.getName() << "\n");
-
-      // Example: Iterate through instructions
-      for (auto &MBB : MF) {
-        for (auto &MI : MBB) {
-          // Process instructions here
-          LLVM_DEBUG(dbgs() << "Instruction: " << MI << "\n");
+    LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
+
+    // Get jump table information
+    MachineJumpTableInfo *JumpTableInfo = MF.getJumpTableInfo();
+    if (!JumpTableInfo) {
+      LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+      return false;
+    }
+    // Assuming JumpTableInfo is available
+    for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
+  const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+  
+  LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " Base Address: " << JTEntry.BaseAddress << "\n");
+
+  // Iterate through the entries (target basic blocks) in this jump table
+  for (auto *MBB : JTEntry.MBBs) {
+    LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << " Address: " << MBB->getAddress() << "\n");
+     }
+    
+
+      // Trace potential indirect jumps related to this jump table
+      traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+    }
+    return false;
+
+}
+
+  void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+    const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+    for (auto &MBB : MF) {
+      for (auto &MI : MBB) {
+        if (MI.isIndirectBranch()) {
+          LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+
+          // Analyze data flow to check if this jump is related to the jump table
+          if (isJumpTableRelated(MI, JTEntry, MF)) {
+            LLVM_DEBUG(dbgs() << "This indirect jump is related to Jump Table #" << JTIndex << "\n");
+          }
         }
       }
+    }
+  }
+
+  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
+  for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
+    const MachineOperand &Op = MI.getOperand(OpIdx);
 
-      return false; // Return true if the pass modifies the function
+    if (Op.isReg()) {
+      Register Reg = Op.getReg();
+      MachineRegisterInfo &MRI = MF.getRegInfo();
+      // Check if any of the definitions of the register are related to a jump table load
+      for (MachineInstr &DefMI : MRI.def_instructions(Reg)) {
+        if (isJumpTableLoad(DefMI, JTEntry)) {
+          return true;
+        }
+      }
+    } else if (Op.isImm()) {
+      // Check if the immediate operand might be an offset/index into the jump table
+      int64_t ImmValue = Op.getImm();
+      
+      // For example, if the jump table has 10 entries, check if the immediate is between 0 and 9
+      if (ImmValue >= 0 && ImmValue < JTEntry.MBBs.size()) {
+        // This immediate value could be an index into the jump table
+        LLVM_DEBUG(dbgs() << "Immediate operand is a possible jump table index: " << ImmValue << "\n");
+        return true;
+      }
     }
+  }
+  return false;
+}
 
-    StringRef getPassName() const override {
-      return "X86 My Backend Pass";
+bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+    if (MI.mayLoad()) {
+        for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+            const MachineOperand &Op = MI.getOperand(i);
+            if (Op.isGlobal() && Op.getGlobal() == JTEntry.BaseAddress) {
+                return true;
+            }
+        }
     }
+    return false;
+}
+
+  StringRef getPassName() const override {
+    return "Match Jump Table Pass";
+  }
+
+    // StringRef getPassName() const override {
+    //   return "X86 My Backend Pass";
+    // }
   };
 }
 
@@ -49,3 +123,7 @@ namespace llvm {
   }
 
 } // end llvm namespace
+
+static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+
+

llvm/lib/Target/X86/X86MatchJumptablePass.h

@@ -1,14 +1,14 @@
 #ifndef LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
 #define LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
 
-
-#pragma once
 #include "llvm/CodeGen/MachineFunctionPass.h"
 
 namespace llvm {
 
-  FunctionPass *createX86MatchJumptablePass();
+/// \brief Creates the X86MatchJumptablePass.
+/// This pass analyzes and processes jump tables in X86 backend code generation.
+FunctionPass *createX86MatchJumptablePass();
 
-} // end namespace llvm
+} // namespace llvm
 
-#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H

llvm/lib/Target/X86/X86TargetMachine.cpp

@@ -11,8 +11,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "X86TargetMachine.h"
+#include "X86MatchJumptablePass.h"
 #include "MCTargetDesc/X86MCTargetDesc.h"
-#include "X86MatchJumptablePass.h"  // Include the header with the function declaration
 #include "TargetInfo/X86TargetInfo.h"
 #include "X86.h"
 #include "X86MachineFunctionInfo.h"
@@ -465,7 +465,7 @@ MachineFunctionInfo *X86TargetMachine::createMachineFunctionInfo(
 
 void X86PassConfig::addIRPasses() {
   addPass(createAtomicExpandLegacyPass());
-  addPass(llvm::createX86MatchJumptablePass());
+
   // We add both pass anyway and when these two passes run, we skip the pass
   // based on the option level and option attribute.
   addPass(createX86LowerAMXIntrinsicsPass());
@@ -597,6 +597,7 @@ void X86PassConfig::addPreEmitPass() {
     addPass(createBreakFalseDeps());
   }
 
+  addPass(createX86MatchJumptablePass());
   addPass(createX86IndirectBranchTrackingPass());
 
   addPass(createX86IssueVZeroUpperPass());