Discover jump table by calling MachineOperand::isJTI()

Author: mingmingl-llvm

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -997,12 +997,6 @@ class MachineBasicBlock
   /// no changes occurred in the meantime.
   bool canSplitCriticalEdge(const MachineBasicBlock *Succ) const;
 
-  /// Return an index for MachineJumpTableInfo if \p this basic block ends with
-  /// an indirect jump using a jump table, otherwise -1.
-  /// This function is a thin wrapper and forward calls to the per-target method
-  /// `TargetInstrInfo::getjumpTableIndex`.
-  int getJumpTableIndex() const;
-
   void pop_front() { Insts.pop_front(); }
   void pop_back() { Insts.pop_back(); }
   void push_back(MachineInstr *MI) { Insts.push_back(MI); }

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -1426,10 +1426,6 @@ bool MachineBasicBlock::canSplitCriticalEdge(
   return true;
 }
 
-int MachineBasicBlock::getJumpTableIndex() const {
-  return findJumpTableIndex(*this);
-}
-
 /// Prepare MI to be removed from its bundle. This fixes bundle flags on MI's
 /// neighboring instructions so the bundle won't be broken by removing MI.
 static void unbundleSingleMI(MachineInstr *MI) {

llvm/lib/CodeGen/StaticDataSplitter.cpp

@@ -92,23 +92,31 @@ bool StaticDataSplitter::splitJumpTablesWithProfiles(
     MachineFunction &MF, MachineJumpTableInfo &MJTI) {
   int NumChangedJumpTables = 0;
 
+  // Jump table could be used by either terminating instructions or
+  // non-terminating ones, so we walk all instructions and use
+  // `MachineOperand::isJTI()` to identify jump table operands.
+  // Similarly, `MachineOperand::isCPI()` can identify constant pool usages
+  // in the same loop.
   for (const auto &MBB : MF) {
-    // IMPORTANT, `getJumpTableIndex` is a thin wrapper around per-target
-    // interface `TargetInstrInfo::getjumpTableIndex`, and only X86 implements
-    // it so far.
-    const int JTI = MBB.getJumpTableIndex();
-    // This is not a source block of jump table.
-    if (JTI == -1)
-      continue;
-
-    auto Hotness = MachineFunctionDataHotness::Hot;
-
-    // Hotness is based on source basic block hotness.
-    if (PSI->isColdBlock(&MBB, MBFI))
-      Hotness = MachineFunctionDataHotness::Cold;
-
-    if (MF.getJumpTableInfo()->updateJumpTableEntryHotness(JTI, Hotness))
-      ++NumChangedJumpTables;
+    for (const MachineInstr &I : MBB) {
+      for (const MachineOperand &Op : I.operands()) {
+        if (!Op.isJTI())
+          continue;
+        const int JTI = Op.getIndex();
+        // This is not a source block of jump table.
+        if (JTI == -1)
+          continue;
+
+        auto Hotness = MachineFunctionDataHotness::Hot;
+
+        // Hotness is based on source basic block hotness.
+        if (PSI->isColdBlock(&MBB, MBFI))
+          Hotness = MachineFunctionDataHotness::Cold;
+
+        if (MF.getJumpTableInfo()->updateJumpTableEntryHotness(JTI, Hotness))
+          ++NumChangedJumpTables;
+      }
+    }
   }
   return NumChangedJumpTables > 0;
 }