[llvm] machinebackend pass test (PR #122487)

Kun Liu via llvm-commits llvm-commits at lists.llvm.org
Sun Jan 12 16:40:51 PST 2025


https://github.com/Ryan-hub-bit updated https://github.com/llvm/llvm-project/pull/122487

>From 9e95cde15ce30630e7118fe046b6dff2d7adf7d9 Mon Sep 17 00:00:00 2001
From: Liu <kliu14 at tulane.edu>
Date: Fri, 10 Jan 2025 10:21:22 -0600
Subject: [PATCH 1/5] machinebackend pass test

---
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 44 +++++++++++++++++++
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  0
 2 files changed, 44 insertions(+)
 create mode 100644 llvm/lib/Target/X86/X86MatchJumptablePass.cpp
 create mode 100644 llvm/lib/Target/X86/X86MatchJumptablePass.h

diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
new file mode 100644
index 00000000000000..bfb6b1b7f89f18
--- /dev/null
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -0,0 +1,44 @@
+#include "X86.h"
+#include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+
+
+using namespace llvm;
+
+namespace {
+  class X86MatchJumptablePass : public MachineFunctionPass {
+  public:
+    static char ID;
+
+    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");
+        }
+      }
+
+      return false; // Return true if the pass modifies the function
+    }
+
+    StringRef getPassName() const override {
+      return "X86 My Backend Pass";
+    }
+  };
+}
+
+char X86MatchJumptablePass::ID = 0;
+
+// Register the pass
+FunctionPass *llvm::createX86MatchJumptablePass() {
+  return new X86MatchJumptablePass();
+}
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
new file mode 100644
index 00000000000000..e69de29bb2d1d6

>From 27bee8dbc8fa03e331302e799429b7450c2b1166 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Fri, 10 Jan 2025 17:01:39 -0600
Subject: [PATCH 2/5] good

---
 .../llvm/Transforms/IPO/JumpTableFinder.h     | 62 +++++++++++++++++++
 llvm/lib/Target/X86/CMakeLists.txt            |  1 +
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 15 +++--
 llvm/lib/Target/X86/X86MatchJumptablePass.h   | 14 +++++
 llvm/lib/Target/X86/X86TargetMachine.cpp      |  3 +-
 5 files changed, 90 insertions(+), 5 deletions(-)
 create mode 100644 llvm/include/llvm/Transforms/IPO/JumpTableFinder.h

diff --git a/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h b/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h
new file mode 100644
index 00000000000000..9d2fb3b405f8c8
--- /dev/null
+++ b/llvm/include/llvm/Transforms/IPO/JumpTableFinder.h
@@ -0,0 +1,62 @@
+// #ifndef LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+// #define LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+
+// #include "llvm/IR/PassManager.h"
+// #include "llvm/IR/Module.h"
+// #include "llvm/Support/raw_ostream.h"
+// #include <set>
+
+// namespace llvm {
+
+// class JumptableFinderPass : public PassInfoMixin<JumptableFinderPass> {
+// public:
+//     /// Main entry point for the pass. Analyzes the module to find and analyze jump tables.
+//     PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+//     /// Implementation of the jump table finder.
+//     void jumptableFinderImpl(Module &M);
+
+//     /// Analyze a SwitchInst for potential jump table patterns.
+//     void findJumpTableFromSwitch(SwitchInst *SI);
+
+//     /// Analyze a GetElementPtrInst for jump table patterns.
+//     void analyzeJumpTable(GetElementPtrInst *GEP);
+
+//     /// Analyze the index computation of a jump table.
+//     void analyzeIndex(Value *Index);
+
+//     /// Find all potential targets for a jump table.
+//     void findTargets(GetElementPtrInst *GEP, std::set<BasicBlock*> &Targets);
+
+//     /// Check the density of a SwitchInst's cases to determine if it forms a jump table.
+//     bool checkDensity(SwitchInst *SI);
+
+//     /// Check if a GetElementPtrInst leads to an indirect branch.
+//     bool leadsToIndirectBranch(GetElementPtrInst *GEP);
+// };
+
+// } // namespace llvm
+
+// #endif // LLVM_ANALYSIS_JUMPTABLEFINDERPASS_H
+
+#ifndef LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
+#define LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
+
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h" // For PassInfoMixin and PreservedAnalyses
+#include "llvm/IR/Function.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/Support/raw_ostream.h"
+#include <set>
+
+namespace llvm {
+
+class JumptableFinderPass : public PassInfoMixin<JumptableFinderPass> {
+public:
+    // Entry point for the pass
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_IPO_JUMPTABLEFINDER_H
diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 9553a8619feb51..4be4576a6a9057 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/llvm/lib/Target/X86/CMakeLists.txt
@@ -23,6 +23,7 @@ tablegen(LLVM X86GenFoldTables.inc -gen-x86-fold-tables -asmwriternum=1)
 add_public_tablegen_target(X86CommonTableGen)
 
 set(sources
+X86MatchJumptablePass.cpp
   X86ArgumentStackSlotRebase.cpp
   X86AsmPrinter.cpp
   X86AvoidTrailingCall.cpp
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index bfb6b1b7f89f18..c3683dacbdf8fe 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -4,7 +4,9 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
+#include "X86MatchJumptablePass.h"
 
+#define DEBUG_TYPE "x86-my-pass"
 
 using namespace llvm;
 
@@ -38,7 +40,12 @@ namespace {
 
 char X86MatchJumptablePass::ID = 0;
 
-// Register the pass
-FunctionPass *llvm::createX86MatchJumptablePass() {
-  return new X86MatchJumptablePass();
-}
+// Ensure the function is in the llvm namespace
+namespace llvm {
+  
+  // Define the pass
+  FunctionPass *createX86MatchJumptablePass() {
+    return new X86MatchJumptablePass();
+  }
+
+} // end llvm namespace
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index e69de29bb2d1d6..67233bf729aec6 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -0,0 +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();
+
+} // end namespace llvm
+
+#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
\ No newline at end of file
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index 20dfdd27b33df6..b9c13fda07dfc7 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/llvm/lib/Target/X86/X86TargetMachine.cpp
@@ -12,6 +12,7 @@
 
 #include "X86TargetMachine.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"
@@ -464,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());

>From 0dc2ceac51b6c4182a8816b215146b3b7a68c1bd Mon Sep 17 00:00:00 2001
From: Liu <kliu14 at tulane.edu>
Date: Fri, 10 Jan 2025 20:05:42 -0600
Subject: [PATCH 3/5] jumptable

---
 llvm/lib/Target/X86/CMakeLists.txt            |   2 +-
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 102 +++++++++++++++---
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  10 +-
 llvm/lib/Target/X86/X86TargetMachine.cpp      |   5 +-
 4 files changed, 99 insertions(+), 20 deletions(-)

diff --git a/llvm/lib/Target/X86/CMakeLists.txt b/llvm/lib/Target/X86/CMakeLists.txt
index 4be4576a6a9057..f5dcf13e47ed25 100644
--- a/llvm/lib/Target/X86/CMakeLists.txt
+++ b/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
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index c3683dacbdf8fe..d8263af59a855d 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/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);
+
+
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index 67233bf729aec6..ac6fc315b025ad 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/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
\ No newline at end of file
+#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
diff --git a/llvm/lib/Target/X86/X86TargetMachine.cpp b/llvm/lib/Target/X86/X86TargetMachine.cpp
index b9c13fda07dfc7..7faeeb2af64e8e 100644
--- a/llvm/lib/Target/X86/X86TargetMachine.cpp
+++ b/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());

>From 65c7421be86c02e11d1e0e6d51b556a22063ba94 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Fri, 10 Jan 2025 21:50:10 -0600
Subject: [PATCH 4/5] test done

---
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 307 +++++++++++++-----
 1 file changed, 234 insertions(+), 73 deletions(-)

diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index d8263af59a855d..c43d7c5b676b49 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -1,23 +1,153 @@
+// #include "X86.h"
+// #include "llvm/CodeGen/MachineFunctionPass.h"
+// #include "llvm/CodeGen/MachineInstr.h"
+// #include "llvm/CodeGen/MachineRegisterInfo.h"
+// #include "llvm/Support/Debug.h"
+// #include "llvm/Support/raw_ostream.h"
+// #include "X86MatchJumptablePass.h"
+// #include "llvm/CodeGen/MachineJumpTableInfo.h"
+
+// #define DEBUG_TYPE "match-jump-table"
+
+// using namespace llvm;
+
+// namespace {
+//   class X86MatchJumptablePass : public MachineFunctionPass {
+//   public:
+//     static char ID;
+
+//     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+
+//     bool runOnMachineFunction(MachineFunction &MF) override {
+//     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);
+
+//     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;
+// }
+
+// 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";
+//     // }
+//   };
+// }
+
+// char X86MatchJumptablePass::ID = 0;
+
+// // Ensure the function is in the llvm namespace
+// namespace llvm {
+  
+//   // Define the pass
+//   FunctionPass *createX86MatchJumptablePass() {
+//     return new X86MatchJumptablePass();
+//   }
+
+// } // end llvm namespace
+
+// static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+
+
 #include "X86.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
-#include "X86MatchJumptablePass.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 
 #define DEBUG_TYPE "match-jump-table"
 
 using namespace llvm;
 
 namespace {
-  class X86MatchJumptablePass : public MachineFunctionPass {
-  public:
-    static char ID;
+class X86MatchJumptablePass : public MachineFunctionPass {
+public:
+  static char ID;
 
-    X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
 
-    bool runOnMachineFunction(MachineFunction &MF) override {
+  bool runOnMachineFunction(MachineFunction &MF) override {
     LLVM_DEBUG(dbgs() << "Analyzing jump tables in function: " << MF.getName() << "\n");
 
     // Get jump table information
@@ -26,104 +156,135 @@ namespace {
       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");
+      const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
-  // 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");
-     }
-    
+      LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " << JTEntry.MBBs.size()
+                        << " entries.\n");
+
+      // Iterate through the entries (target basic blocks) in this jump table
+      for (auto *MBB : JTEntry.MBBs) {
+        if (MBB) {
+          LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << "\n");
+        }
+      }
 
-      // Trace potential indirect jumps related to this jump table
-      traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+      // Assuming you have access to MF, JTIndex, and JumpTableInfo
+    MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+
+      if (indirectJumpInstr) {
+          // Handle the found indirect jump instruction
+          dbgs() << "Indirect jump found at address: " << indirectJumpInstr << "\n";
+          for (auto &MBB : JTEntry.MBBs) {
+          LLVM_DEBUG(dbgs() << "Address of MBB: " << &MBB << "\n"); // Print the address of the MBB
+          // Optionally print the name and instructions inside the MBB
+         // LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
+          // for (auto &MI : MBB) {
+          //     LLVM_DEBUG(dbgs() << "  Instruction: " << MI << "\n");
+          // }
+        }
+      } 
     }
+
     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");
 
-  void traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+  //         // 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");
+  //         }
+  //       }
+  //     }
+  //   }
+  // }
+  MachineInstr* 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");
+        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");
-          }
+                // 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");
+                    
+                    // Return the address of the indirect jump (MI)
+                    return &MI;
+                }
+            }
         }
-      }
     }
-  }
+
+    // Return nullptr if no indirect jump is found
+    return nullptr;
+}
 
   bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
-  for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
-    const MachineOperand &Op = MI.getOperand(OpIdx);
-
-    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)) {
+    for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
+      const MachineOperand &Op = MI.getOperand(OpIdx);
+
+      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();
+        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;
         }
       }
-    } 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;
   }
-  return false;
-}
 
-bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+  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;
-            }
+      for (unsigned i = 0; i < MI.getNumOperands(); ++i) {
+        const MachineOperand &Op = MI.getOperand(i);
+          if (Op.getType() == MachineOperand::MO_JumpTableIndex) {
+          unsigned JTIndex = Op.getIndex();
+          if (JTIndex < JTEntry.MBBs.size()) {
+            LLVM_DEBUG(dbgs() << "Instruction loads from Jump Table index: " << JTIndex << "\n");
+            return true;
+          }
         }
+      }
     }
     return false;
-}
-
-  StringRef getPassName() const override {
-    return "Match Jump Table Pass";
   }
 
-    // StringRef getPassName() const override {
-    //   return "X86 My Backend Pass";
-    // }
-  };
-}
+  StringRef getPassName() const override { return "Match Jump Table Pass"; }
+};
+
+} // namespace
 
 char X86MatchJumptablePass::ID = 0;
 
-// Ensure the function is in the llvm namespace
 namespace llvm {
-  
-  // Define the pass
-  FunctionPass *createX86MatchJumptablePass() {
-    return new X86MatchJumptablePass();
-  }
-
-} // end llvm namespace
 
-static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+// Define the pass
+FunctionPass *createX86MatchJumptablePass() { return new X86MatchJumptablePass(); }
 
+} // namespace llvm
 
+static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false,
+                                             false);

>From 9dd981b98af2b53d024ea7a10064b49b8daa99f9 Mon Sep 17 00:00:00 2001
From: ryan <liu864579887 at gmail.com>
Date: Sun, 12 Jan 2025 18:40:37 -0600
Subject: [PATCH 5/5] clang jump table pattern test successful

---
 llvm/lib/Target/X86/X86.h                     |   2 +
 llvm/lib/Target/X86/X86MatchJumptablePass.cpp | 504 +++++++++---------
 llvm/lib/Target/X86/X86MatchJumptablePass.h   |  46 +-
 3 files changed, 308 insertions(+), 244 deletions(-)

diff --git a/llvm/lib/Target/X86/X86.h b/llvm/lib/Target/X86/X86.h
index 48a3fe1934a967..5d4882b921a23c 100644
--- a/llvm/lib/Target/X86/X86.h
+++ b/llvm/lib/Target/X86/X86.h
@@ -54,6 +54,8 @@ FunctionPass *createX86IndirectBranchTrackingPass();
 /// This will prevent a stall when returning on the Atom.
 FunctionPass *createX86PadShortFunctions();
 
+FunctionPass *createX86MatchJumptablePass();
+
 /// Return a pass that selectively replaces certain instructions (like add,
 /// sub, inc, dec, some shifts, and some multiplies) by equivalent LEA
 /// instructions, in order to eliminate execution delays in some processors.
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
index c43d7c5b676b49..cf2a8ffe56345f 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.cpp
@@ -1,290 +1,316 @@
-// #include "X86.h"
-// #include "llvm/CodeGen/MachineFunctionPass.h"
-// #include "llvm/CodeGen/MachineInstr.h"
-// #include "llvm/CodeGen/MachineRegisterInfo.h"
-// #include "llvm/Support/Debug.h"
-// #include "llvm/Support/raw_ostream.h"
-// #include "X86MatchJumptablePass.h"
-// #include "llvm/CodeGen/MachineJumpTableInfo.h"
-
-// #define DEBUG_TYPE "match-jump-table"
-
-// using namespace llvm;
-
-// namespace {
-//   class X86MatchJumptablePass : public MachineFunctionPass {
-//   public:
-//     static char ID;
-
-//     X86MatchJumptablePass() : MachineFunctionPass(ID) {}
-
-//     bool runOnMachineFunction(MachineFunction &MF) override {
-//     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);
-
-//     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;
-// }
-
-// 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";
-//     // }
-//   };
-// }
-
-// char X86MatchJumptablePass::ID = 0;
-
-// // Ensure the function is in the llvm namespace
-// namespace llvm {
-  
-//   // Define the pass
-//   FunctionPass *createX86MatchJumptablePass() {
-//     return new X86MatchJumptablePass();
-//   }
-
-// } // end llvm namespace
-
-// static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
-
-
+// X86MatchJumptablePass.cpp
+#include "X86MatchJumptablePass.h"
 #include "X86.h"
+#include "X86InstrInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
-#include "llvm/CodeGen/MachineInstr.h"
-#include "llvm/CodeGen/MachineRegisterInfo.h"
-#include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
+#include "llvm/PassRegistry.h"
+#include "llvm/Pass.h"
 
 #define DEBUG_TYPE "match-jump-table"
 
-using namespace llvm;
+namespace llvm {
 
-namespace {
-class X86MatchJumptablePass : public MachineFunctionPass {
-public:
-  static char ID;
+char X86MatchJumptablePass::ID = 0;
 
-  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+void initializeX86MatchJumptablePassPass(PassRegistry &Registry) {
+  RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false, false);
+}
 
-  bool runOnMachineFunction(MachineFunction &MF) override {
+void X86MatchJumptablePass::insertIdentifyingMarker(MachineInstr* MI, 
+                                                   MachineFunction &MF, 
+                                                   unsigned JTIndex) {
+  // Create a unique identifier for this indirect jump
+  std::string marker = "IJUMP_" + MF.getName().str() + "_" + std::to_string(JTIndex);
+  
+  // Get target instruction info
+  const TargetInstrInfo *TII = MF.getSubtarget().getInstrInfo();
+  
+  // Insert a special no-op instruction before the indirect jump with debug info
+  BuildMI(*MI->getParent(), MI, MI->getDebugLoc(), 
+          TII->get(X86::NOOP))
+          .addExternalSymbol(marker.c_str());
+}
+
+void X86MatchJumptablePass::recordJumpLocation(MachineInstr* MI, 
+                                              MachineFunction &MF, 
+                                              unsigned JTIndex) {
+  // Create a new section entry
+  auto &Context = MF.getFunction().getContext();
+  auto *M = MF.getFunction().getParent();
+  
+  // Create mapping data
+  std::string sectionName = ".jumptable.map";
+  std::string data = MF.getName().str() + "," + 
+                    std::to_string(JTIndex) + "\n";
+  
+  // Add the data to a special section
+  auto *GV = new GlobalVariable(
+    *M,
+    ArrayType::get(Type::getInt8Ty(Context), data.size() + 1),
+    true,
+    GlobalValue::PrivateLinkage,
+    ConstantDataArray::getString(Context, data),
+    "jump_map_entry"
+  );
+  
+  GV->setSection(sectionName);
+}
+
+bool X86MatchJumptablePass::runOnMachineFunction(MachineFunction &MF) {
+    static int PassRunCount = 0;
+    PassRunCount++;
+    
+    LLVM_DEBUG(dbgs() << "\n=== Pass Run #" << PassRunCount << " ===\n");
+    LLVM_DEBUG(dbgs() << "Current optimization phase: " 
+               << MF.getFunction().getParent()->getSourceFileName() << "\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;
+        LLVM_DEBUG(dbgs() << "No jump tables in this function.\n");
+        return false;
     }
 
+    bool Modified = false;
     for (unsigned JTIndex = 0; JTIndex < JumpTableInfo->getJumpTables().size(); ++JTIndex) {
-      const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+        const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
+
+        LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " 
+                         << JTEntry.MBBs.size() << " entries.\n");
+
+        // Print detailed information about jump table entries
+        LLVM_DEBUG(dbgs() << "Jump table entries:\n");
+        for (unsigned i = 0; i < JTEntry.MBBs.size(); ++i) {
+            if (JTEntry.MBBs[i]) {
+                LLVM_DEBUG(dbgs() << "  [" << i << "] -> BB#" 
+                           << JTEntry.MBBs[i]->getNumber() 
+                           << " (" << JTEntry.MBBs[i]->getName() << ")\n");
+            }
+        }
 
-      LLVM_DEBUG(dbgs() << "Jump Table #" << JTIndex << " contains " << JTEntry.MBBs.size()
-                        << " entries.\n");
+        // Find and process indirect jumps
+        MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
+
+        if (indirectJumpInstr) {
+            // Print more detailed information about the indirect jump
+            LLVM_DEBUG(dbgs() << "Found indirect jump: " << *indirectJumpInstr << "\n");
+            LLVM_DEBUG(dbgs() << "Parent basic block: " << indirectJumpInstr->getParent()->getName() << "\n");
+            
+            // Print operand information
+            LLVM_DEBUG(dbgs() << "Jump instruction operands:\n");
+            for (unsigned i = 0; i < indirectJumpInstr->getNumOperands(); ++i) {
+                LLVM_DEBUG(dbgs() << "  Operand " << i << ": " << indirectJumpInstr->getOperand(i) << "\n");
+            }
 
-      // Iterate through the entries (target basic blocks) in this jump table
-      for (auto *MBB : JTEntry.MBBs) {
-        if (MBB) {
-          LLVM_DEBUG(dbgs() << "  Target BasicBlock: " << MBB->getName() << "\n");
-        }
-      }
-
-      // Assuming you have access to MF, JTIndex, and JumpTableInfo
-    MachineInstr* indirectJumpInstr = traceIndirectJumps(MF, JTIndex, JumpTableInfo);
-
-      if (indirectJumpInstr) {
-          // Handle the found indirect jump instruction
-          dbgs() << "Indirect jump found at address: " << indirectJumpInstr << "\n";
-          for (auto &MBB : JTEntry.MBBs) {
-          LLVM_DEBUG(dbgs() << "Address of MBB: " << &MBB << "\n"); // Print the address of the MBB
-          // Optionally print the name and instructions inside the MBB
-         // LLVM_DEBUG(dbgs() << "MBB: " << MBB.getName() << "\n");
-          // for (auto &MI : MBB) {
-          //     LLVM_DEBUG(dbgs() << "  Instruction: " << MI << "\n");
-          // }
+            // Insert marker and record location
+            insertIdentifyingMarker(indirectJumpInstr, MF, JTIndex);
+            recordJumpLocation(indirectJumpInstr, MF, JTIndex);
+            Modified = true;
         }
-      } 
     }
 
-    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");
-  //         }
-  //       }
-  //     }
-  //   }
-  // }
-  MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, MachineJumpTableInfo *JumpTableInfo) {
+    return Modified;
+}
+
+MachineInstr* X86MatchJumptablePass::traceIndirectJumps(MachineFunction &MF, 
+                                                        unsigned JTIndex,
+                                                        MachineJumpTableInfo *JumpTableInfo) {
     const MachineJumpTableEntry &JTEntry = JumpTableInfo->getJumpTables()[JTIndex];
 
+    LLVM_DEBUG(dbgs() << "Tracing indirect jumps:\n");
     for (auto &MBB : MF) {
+        LLVM_DEBUG(dbgs() << "  Checking BB: " << MBB.getName() << "\n");
         for (auto &MI : MBB) {
+            LLVM_DEBUG(dbgs() << "    Checking instruction: " << MI << "\n");
             if (MI.isIndirectBranch()) {
-                LLVM_DEBUG(dbgs() << "Found indirect jump: " << MI << "\n");
+                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");
-                    
-                    // Return the address of the indirect jump (MI)
+                    LLVM_DEBUG(dbgs() << "    This indirect jump is related to Jump Table #" 
+                              << JTIndex << "\n");
                     return &MI;
+                } else {
+                    LLVM_DEBUG(dbgs() << "    Jump is not related to this jump table\n");
                 }
             }
         }
     }
 
-    // Return nullptr if no indirect jump is found
+    LLVM_DEBUG(dbgs() << "  No related indirect jump found\n");
     return nullptr;
 }
 
-  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, MachineFunction &MF) {
-    for (unsigned OpIdx = 0; OpIdx < MI.getNumOperands(); ++OpIdx) {
-      const MachineOperand &Op = MI.getOperand(OpIdx);
+bool X86MatchJumptablePass::isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry) {
+    LLVM_DEBUG(dbgs() << "\nAnalyzing potential jump table load instruction: " << MI << "\n");
+
+    // First check memory operands for jump table metadata
+    for (const MachineMemOperand *MMO : MI.memoperands()) {
+        LLVM_DEBUG(dbgs() << "  Checking memory operand flags: " << MMO->getFlags() << "\n");
+        if (MMO->getValue()) {
+            StringRef ValueName = MMO->getValue()->getName();
+            LLVM_DEBUG(dbgs() << "    Memory value name: '" << ValueName << "'\n");
+            if (ValueName.contains("jump-table")) {
+                LLVM_DEBUG(dbgs() << "    Found jump table in memory value name\n");
+                return true;
+            }
+        }
 
-      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;
-          }
+        // Check if this is a jump table load directly from memory operand comments
+        if (MI.getDesc().mayLoad() && MI.hasOneMemOperand()) {
+            // Look for jump table reference in the instruction's debug info or comments
+            if (MI.getDebugLoc()) {
+                std::string Comment;
+                raw_string_ostream OS(Comment);
+                MI.print(OS);
+                if (Comment.find("jump-table") != std::string::npos) {
+                    LLVM_DEBUG(dbgs() << "    Found jump table reference in instruction comment\n");
+                    return true;
+                }
+            }
+        }
+    }
+
+    // Check for the MOVSX pattern
+    if (MI.getOpcode() == X86::MOVSX64rm32) {
+        LLVM_DEBUG(dbgs() << "  Found MOVSX64rm32 instruction\n");
+        Register BaseReg;
+        
+        // Find base register
+        for (const MachineOperand &MO : MI.operands()) {
+            if (MO.isReg() && MO.isUse()) {
+                BaseReg = MO.getReg();
+                LLVM_DEBUG(dbgs() << "    Found base register: " << printReg(BaseReg, nullptr) << "\n");
+                break;
+            }
         }
-      } else if (Op.isImm()) {
-        // Check if the immediate operand might be an offset/index into the jump table
-        int64_t ImmValue = Op.getImm();
-        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;
+
+        if (BaseReg) {
+            // Look for preceding LEA
+            MachineBasicBlock::iterator MBBI = MI;
+            const MachineBasicBlock *MBB = MI.getParent();
+            
+            LLVM_DEBUG(dbgs() << "    Looking for LEA defining register: " << printReg(BaseReg, nullptr) << "\n");
+            
+            while (MBBI != MBB->begin()) {
+                --MBBI;
+                LLVM_DEBUG(dbgs() << "      Checking: " << *MBBI << "\n");
+                
+                if (MBBI->getOpcode() == X86::LEA64r) {
+                    LLVM_DEBUG(dbgs() << "      Found LEA64r\n");
+                    
+                    // Verify this LEA defines our base register
+                    const MachineOperand &DefReg = MBBI->getOperand(0);
+                    if (!DefReg.isReg() || DefReg.getReg() != BaseReg) {
+                        LLVM_DEBUG(dbgs() << "      LEA defines different register\n");
+                        continue;
+                    }
+                    
+                    // Check for jump table symbol
+                    for (const MachineOperand &MO : MBBI->operands()) {
+                        if (MO.isSymbol()) {
+                            StringRef SymName = MO.getSymbolName();
+                            LLVM_DEBUG(dbgs() << "      Checking symbol: '" << SymName << "'\n");
+                            if (SymName.contains("jump-table")) {
+                                LLVM_DEBUG(dbgs() << "      Found jump table symbol!\n");
+                                return true;
+                            }
+                        }
+                    }
+                }
+            }
+            LLVM_DEBUG(dbgs() << "    No matching LEA found\n");
         }
-      }
     }
+
     return false;
-  }
-
-  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.getType() == MachineOperand::MO_JumpTableIndex) {
-          unsigned JTIndex = Op.getIndex();
-          if (JTIndex < JTEntry.MBBs.size()) {
-            LLVM_DEBUG(dbgs() << "Instruction loads from Jump Table index: " << JTIndex << "\n");
-            return true;
-          }
+}
+
+bool X86MatchJumptablePass::isJumpTableRelated(MachineInstr &MI, 
+                                              const MachineJumpTableEntry &JTEntry,
+                                              MachineFunction &MF) {
+    if (!MI.isIndirectBranch()) {
+        LLVM_DEBUG(dbgs() << "Not an indirect branch, skipping\n");
+        return false;
+    }
+
+    LLVM_DEBUG(dbgs() << "\nAnalyzing indirect jump: " << MI << "\n");
+
+    // Get jump register
+    Register JumpReg;
+    for (const MachineOperand &MO : MI.operands()) {
+        if (MO.isReg() && MO.isUse()) {
+            JumpReg = MO.getReg();
+            LLVM_DEBUG(dbgs() << "Found jump register: " << printReg(JumpReg, nullptr) << "\n");
+            break;
         }
-      }
     }
-    return false;
-  }
 
-  StringRef getPassName() const override { return "Match Jump Table Pass"; }
-};
+    if (!JumpReg) {
+        LLVM_DEBUG(dbgs() << "No jump register found\n");
+        return false;
+    }
 
-} // namespace
+    SmallVector<MachineInstr*, 8> Worklist;
+    SmallPtrSet<MachineInstr*, 16> Visited;
+    
+    LLVM_DEBUG(dbgs() << "Starting backward analysis from register " << printReg(JumpReg, nullptr) << "\n");
 
-char X86MatchJumptablePass::ID = 0;
+    for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(JumpReg)) {
+        Worklist.push_back(&DefMI);
+        LLVM_DEBUG(dbgs() << "Added to worklist: " << DefMI << "\n");
+    }
 
-namespace llvm {
+    while (!Worklist.empty()) {
+        MachineInstr *CurrMI = Worklist.pop_back_val();
+        if (!Visited.insert(CurrMI).second) {
+            LLVM_DEBUG(dbgs() << "Already visited: " << *CurrMI << "\n");
+            continue;
+        }
+
+        LLVM_DEBUG(dbgs() << "Analyzing instruction: " << *CurrMI << "\n");
+
+        if (isJumpTableLoad(*CurrMI, JTEntry)) {
+            LLVM_DEBUG(dbgs() << "Found jump table load!\n");
+            return true;
+        }
+
+        if (CurrMI->getOpcode() == X86::ADD64rr) {
+            LLVM_DEBUG(dbgs() << "Found ADD64rr, checking operands\n");
+            for (const MachineOperand &MO : CurrMI->operands()) {
+                if (MO.isReg() && MO.isUse()) {
+                    LLVM_DEBUG(dbgs() << "Checking register operand: " << printReg(MO.getReg(), nullptr) << "\n");
+                    for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(MO.getReg())) {
+                        if (isJumpTableLoad(DefMI, JTEntry)) {
+                            LLVM_DEBUG(dbgs() << "Found jump table load via ADD operand!\n");
+                            return true;
+                        }
+                    }
+                }
+            }
+        }
 
-// Define the pass
-FunctionPass *createX86MatchJumptablePass() { return new X86MatchJumptablePass(); }
+        // Add uses to worklist
+        for (const MachineOperand &MO : CurrMI->operands()) {
+            if (MO.isReg() && MO.isUse()) {
+                LLVM_DEBUG(dbgs() << "Adding definitions of register " << printReg(MO.getReg(), nullptr) << " to worklist\n");
+                for (MachineInstr &DefMI : MF.getRegInfo().def_instructions(MO.getReg())) {
+                    if (!Visited.count(&DefMI)) {
+                        Worklist.push_back(&DefMI);
+                        LLVM_DEBUG(dbgs() << "Added to worklist: " << DefMI << "\n");
+                    }
+                }
+            }
+        }
+    }
 
-} // namespace llvm
+    LLVM_DEBUG(dbgs() << "No jump table relation found\n");
+    return false;
+}
+
+FunctionPass *createX86MatchJumptablePass() { 
+    return new X86MatchJumptablePass();
+}
 
-static RegisterPass<X86MatchJumptablePass> X("match-jump-table", "Match Jump Table Pass", false,
-                                             false);
+} // end namespace llvm
\ No newline at end of file
diff --git a/llvm/lib/Target/X86/X86MatchJumptablePass.h b/llvm/lib/Target/X86/X86MatchJumptablePass.h
index ac6fc315b025ad..779ada3d8b6d41 100644
--- a/llvm/lib/Target/X86/X86MatchJumptablePass.h
+++ b/llvm/lib/Target/X86/X86MatchJumptablePass.h
@@ -1,14 +1,50 @@
-#ifndef LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
-#define LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+// X86MatchJumptablePass.h
+#ifndef X86_MATCH_JUMPTABLE_PASS_H
+#define X86_MATCH_JUMPTABLE_PASS_H
 
+#include "MCTargetDesc/X86MCTargetDesc.h"
+#include "X86.h"
+#include "X86InstrInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineRegisterInfo.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/IR/GlobalVariable.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/CodeGen/TargetInstrInfo.h"
 
 namespace llvm {
 
-/// \brief Creates the X86MatchJumptablePass.
-/// This pass analyzes and processes jump tables in X86 backend code generation.
+class X86MatchJumptablePass : public MachineFunctionPass {
+private:
+  void insertIdentifyingMarker(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
+  void recordJumpLocation(MachineInstr* MI, MachineFunction &MF, unsigned JTIndex);
+  MachineInstr* traceIndirectJumps(MachineFunction &MF, unsigned JTIndex, 
+                                  MachineJumpTableInfo *JumpTableInfo);
+  bool isJumpTableRelated(MachineInstr &MI, const MachineJumpTableEntry &JTEntry, 
+                         MachineFunction &MF);
+  bool isJumpTableLoad(MachineInstr &MI, const MachineJumpTableEntry &JTEntry);
+  bool isRegUsedInJumpTableLoad(Register Reg,MachineFunction &MF,
+                                                    const MachineJumpTableEntry &JTEntry);
+
+public:
+  static char ID;
+
+  X86MatchJumptablePass() : MachineFunctionPass(ID) {}
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+  StringRef getPassName() const override { return "Match Jump Table Pass"; }
+};
+
 FunctionPass *createX86MatchJumptablePass();
 
+// Pass initialization declaration
+void initializeX86MatchJumptablePassPass(PassRegistry &Registry);
+
 } // namespace llvm
 
-#endif // LLVM_LIB_TARGET_X86_X86MATCHJUMPTABLEPASS_H
+#endif // X86_MATCH_JUMPTABLE_PASS_H
+



More information about the llvm-commits mailing list