[LoongArch] Add options for annotate tablejump

clang/include/clang/Driver/Options.td

@@ -5384,9 +5384,14 @@ def mlasx : Flag<["-"], "mlasx">, Group<m_loongarch_Features_Group>,
   HelpText<"Enable Loongson Advanced SIMD Extension (LASX).">;
 def mno_lasx : Flag<["-"], "mno-lasx">, Group<m_loongarch_Features_Group>,
   HelpText<"Disable Loongson Advanced SIMD Extension (LASX).">;
+let Flags = [TargetSpecific] in {
 def msimd_EQ : Joined<["-"], "msimd=">, Group<m_loongarch_Features_Group>,
-  Flags<[TargetSpecific]>,
   HelpText<"Select the SIMD extension(s) to be enabled in LoongArch either 'none', 'lsx', 'lasx'.">;
+def mannotate_tablejump : Flag<["-"], "mannotate-tablejump">, Group<m_loongarch_Features_Group>,
+  HelpText<"Enable annotate table jump instruction to correlate it with the jump table.">;
+def mno_annotate_tablejump : Flag<["-"], "mno-annotate-tablejump">, Group<m_loongarch_Features_Group>,
+  HelpText<"Disable annotate table jump instruction to correlate it with the jump table.">;
+} // let Flags = [TargetSpecific]
 def mnop_mcount : Flag<["-"], "mnop-mcount">, HelpText<"Generate mcount/__fentry__ calls as nops. To activate they need to be patched in.">,
   Visibility<[ClangOption, CC1Option]>, Group<m_Group>,
   MarshallingInfoFlag<CodeGenOpts<"MNopMCount">>;

clang/lib/Driver/ToolChains/Clang.cpp

@@ -1870,6 +1870,14 @@ void Clang::AddLoongArchTargetArgs(const ArgList &Args,
     CmdArgs.push_back("-tune-cpu");
     CmdArgs.push_back(Args.MakeArgString(TuneCPU));
   }
+
+  if (Arg *A = Args.getLastArg(options::OPT_mannotate_tablejump,
+                               options::OPT_mno_annotate_tablejump)) {
+    if (A->getOption().matches(options::OPT_mannotate_tablejump)) {
+      CmdArgs.push_back("-mllvm");
+      CmdArgs.push_back("-loongarch-annotate-tablejump");
+    }
+  }
 }
 
 void Clang::AddMIPSTargetArgs(const ArgList &Args,

clang/test/Driver/loongarch-mannotate-tablejump.c

@@ -0,0 +1,13 @@
+/// Test -m[no-]annotate-tablejump options.
+
+// RUN: %clang --target=loongarch64 -mannotate-tablejump %s -### 2>&1 | \
+// RUN:   FileCheck %s --check-prefix=CC1-ANOTATE
+// RUN: %clang --target=loongarch64 -mno-annotate-tablejump %s -### 2>&1 | \
+// RUN:   FileCheck %s --check-prefix=CC1-NO-ANOTATE
+// RUN: %clang --target=loongarch64 -mannotate-tablejump -mno-annotate-tablejump %s -### 2>&1 | \
+// RUN:   FileCheck %s --check-prefix=CC1-NO-ANOTATE
+// RUN: %clang --target=loongarch64 -mno-annotate-tablejump -mannotate-tablejump %s -### 2>&1 | \
+// RUN:   FileCheck %s --check-prefix=CC1-ANOTATE
+
+// CC1-ANOTATE: "-loongarch-annotate-tablejump"
+// CC1-NO-ANOTATE-NOT: "-loongarch-annotate-tablejump"

llvm/lib/Target/LoongArch/LoongArchAsmPrinter.cpp

@@ -13,18 +13,30 @@
 
 #include "LoongArchAsmPrinter.h"
 #include "LoongArch.h"
+#include "LoongArchInstrInfo.h"
+#include "LoongArchMachineFunctionInfo.h"
 #include "LoongArchTargetMachine.h"
 #include "MCTargetDesc/LoongArchInstPrinter.h"
+#include "MCTargetDesc/LoongArchMCTargetDesc.h"
 #include "TargetInfo/LoongArchTargetInfo.h"
 #include "llvm/CodeGen/AsmPrinter.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
+#include "llvm/CodeGen/MachineModuleInfoImpls.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstBuilder.h"
+#include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/TargetRegistry.h"
 
 using namespace llvm;
 
 #define DEBUG_TYPE "loongarch-asm-printer"
 
+cl::opt<bool> LArchAnnotateTableJump(
+    "loongarch-annotate-tablejump", cl::Hidden,
+    cl::desc(
+        "Annotate table jump instruction to correlate it with the jump table."),
+    cl::init(false));
+
 // Simple pseudo-instructions have their lowering (with expansion to real
 // instructions) auto-generated.
 #include "LoongArchGenMCPseudoLowering.inc"
@@ -238,6 +250,39 @@ void LoongArchAsmPrinter::emitSled(const MachineInstr &MI, SledKind Kind) {
   recordSled(BeginOfSled, MI, Kind, 2);
 }
 
+void LoongArchAsmPrinter::emitJumpTableInfo() {
+  AsmPrinter::emitJumpTableInfo();
+
+  if (!LArchAnnotateTableJump)
+    return;
+
+  assert(TM.getTargetTriple().isOSBinFormatELF());
+
+  unsigned Size = getDataLayout().getPointerSize();
+  auto *LAFI = MF->getInfo<LoongArchMachineFunctionInfo>();
+  unsigned EntrySize = LAFI->getJumpInfoSize();
+
+  if (0 == EntrySize)
+    return;
+
+  // Emit an additional section to store the correlation info as pairs of
+  // addresses, each pair contains the address of a jump instruction (jr) and
+  // the address of the jump table.
+  OutStreamer->switchSection(MMI->getContext().getELFSection(
+      ".discard.tablejump_annotate", ELF::SHT_PROGBITS, 0));
+
+  for (unsigned Idx = 0; Idx < EntrySize; ++Idx) {
+    OutStreamer->emitValue(
+        MCSymbolRefExpr::create(LAFI->getJumpInfoJrMI(Idx)->getPreInstrSymbol(),
+                                OutContext),
+        Size);
+    OutStreamer->emitValue(
+        MCSymbolRefExpr::create(
+            GetJTISymbol(LAFI->getJumpInfoJTIMO(Idx)->getIndex()), OutContext),
+        Size);
+  }
+}
+
 bool LoongArchAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
   AsmPrinter::runOnMachineFunction(MF);
   // Emit the XRay table for this function.

llvm/lib/Target/LoongArch/LoongArchAsmPrinter.h

@@ -55,6 +55,7 @@ class LLVM_LIBRARY_VISIBILITY LoongArchAsmPrinter : public AsmPrinter {
   bool lowerOperand(const MachineOperand &MO, MCOperand &MCOp) const {
     return lowerLoongArchMachineOperandToMCOperand(MO, MCOp, *this);
   }
+  void emitJumpTableInfo() override;
 };
 
 } // end namespace llvm

llvm/lib/Target/LoongArch/LoongArchExpandPseudoInsts.cpp

@@ -13,6 +13,7 @@
 
 #include "LoongArch.h"
 #include "LoongArchInstrInfo.h"
+#include "LoongArchMachineFunctionInfo.h"
 #include "LoongArchTargetMachine.h"
 #include "MCTargetDesc/LoongArchBaseInfo.h"
 #include "MCTargetDesc/LoongArchMCTargetDesc.h"
@@ -27,6 +28,8 @@
 
 using namespace llvm;
 
+extern cl::opt<bool> LArchAnnotateTableJump;
+
 #define LOONGARCH_PRERA_EXPAND_PSEUDO_NAME                                     \
   "LoongArch Pre-RA pseudo instruction expansion pass"
 #define LOONGARCH_EXPAND_PSEUDO_NAME                                           \
@@ -103,6 +106,8 @@ class LoongArchPreRAExpandPseudo : public MachineFunctionPass {
                           MachineBasicBlock::iterator MBBI,
                           MachineBasicBlock::iterator &NextMBBI,
                           bool IsTailCall);
+  void annotateTableJump(MachineBasicBlock &MBB,
+                         MachineBasicBlock::iterator MBBI);
 };
 
 char LoongArchPreRAExpandPseudo::ID = 0;
@@ -167,6 +172,12 @@ bool LoongArchPreRAExpandPseudo::expandMI(
   case LoongArch::PseudoTAIL_MEDIUM:
   case LoongArch::PseudoTAIL_LARGE:
     return expandFunctionCALL(MBB, MBBI, NextMBBI, /*IsTailCall=*/true);
+  case LoongArch::PseudoBRIND:
+    // If the PseudoBRIND is used to table jump, then emit a label to annotate
+    // the `jr` instruction, and save the instructions.
+    if (LArchAnnotateTableJump)
+      annotateTableJump(MBB, MBBI);
+    break;
   }
   return false;
 }
@@ -601,6 +612,44 @@ bool LoongArchPreRAExpandPseudo::expandFunctionCALL(
   return true;
 }
 
+void LoongArchPreRAExpandPseudo::annotateTableJump(
+    MachineBasicBlock &MBB, MachineBasicBlock::iterator MBBI) {
+  MachineFunction *MF = MBB.getParent();
+  MachineRegisterInfo &MRI = MBB.getParent()->getRegInfo();
+
+  bool IsFound = false;
+
+  std::function<void(MachineInstr *, int)> FindJTIMI = [&](MachineInstr *MInst,
+                                                           int FindDepth) {
+    if (FindDepth < 0)
+      return;
+    for (auto &MO : MInst->all_uses()) {
+      if (IsFound)
+        return;
+      Register Reg = MO.getReg();
+      if (!Reg.isVirtual())
+        continue;
+      MachineInstr *DefMI = MRI.getVRegDef(Reg);
+      if (!DefMI)
+        continue;
+      for (unsigned Idx = 0; Idx < DefMI->getNumOperands(); ++Idx) {
+        MachineOperand &MO = DefMI->getOperand(Idx);
+        if (MO.isJTI()) {
+          MBBI->setPreInstrSymbol(
+              *MF, MF->getContext().createNamedTempSymbol("jrtb_"));
+          MF->getInfo<LoongArchMachineFunctionInfo>()->setJumpInfo(&*MBBI, &MO);
+          IsFound = true;
+          return;
+        }
+      }
+      FindJTIMI(DefMI, --FindDepth);
+    }
+  };
+
+  // FindDepth = 3, probably sufficient.
+  FindJTIMI(&*MBBI, /*FindDepth=*/3);
+}
+
 class LoongArchExpandPseudo : public MachineFunctionPass {
 public:
   const LoongArchInstrInfo *TII;

llvm/lib/Target/LoongArch/LoongArchMachineFunctionInfo.h

@@ -39,6 +39,10 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
   /// Registers that have been sign extended from i32.
   SmallVector<Register, 8> SExt32Registers;
 
+  /// Pairs of `jr` instructions and corresponding JTI operands, used for the
+  /// `annotate-tablejump` option.
+  SmallVector<std::pair<MachineInstr *, MachineOperand *>, 4> JumpInfos;
+
 public:
   LoongArchMachineFunctionInfo(const Function &F,
                                const TargetSubtargetInfo *STI) {}
@@ -71,6 +75,15 @@ class LoongArchMachineFunctionInfo : public MachineFunctionInfo {
   bool isSExt32Register(Register Reg) const {
     return is_contained(SExt32Registers, Reg);
   }
+
+  void setJumpInfo(MachineInstr *JrMI, MachineOperand *JTIMO) {
+    JumpInfos.push_back(std::make_pair(JrMI, JTIMO));
+  }
+  unsigned getJumpInfoSize() { return JumpInfos.size(); }
+  MachineInstr *getJumpInfoJrMI(unsigned Idx) { return JumpInfos[Idx].first; }
+  MachineOperand *getJumpInfoJTIMO(unsigned Idx) {
+    return JumpInfos[Idx].second;
+  }
 };
 
 } // end namespace llvm

llvm/test/CodeGen/LoongArch/annotate-tablejump.ll

@@ -0,0 +1,133 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 5
+; RUN: llc --mtriple=loongarch32 -mattr=+d \
+; RUN:   --min-jump-table-entries=4 < %s \
+; RUN:   --loongarch-annotate-tablejump \
+; RUN:   | FileCheck %s --check-prefix=LA32-JT
+; RUN: llc --mtriple=loongarch64 -mattr=+d \
+; RUN:   --min-jump-table-entries=4 < %s \
+; RUN:   --loongarch-annotate-tablejump \
+; RUN:   | FileCheck %s --check-prefix=LA64-JT
+
+define void @switch_4_arms(i32 %in, ptr %out) nounwind {
+; LA32-JT-LABEL: switch_4_arms:
+; LA32-JT:       # %bb.0: # %entry
+; LA32-JT-NEXT:    addi.w $a3, $a0, -1
+; LA32-JT-NEXT:    ori $a2, $zero, 3
+; LA32-JT-NEXT:    bltu $a2, $a3, .LBB0_7
+; LA32-JT-NEXT:  # %bb.1: # %entry
+; LA32-JT-NEXT:    pcalau12i $a4, %pc_hi20(.LJTI0_0)
+; LA32-JT-NEXT:    addi.w $a4, $a4, %pc_lo12(.LJTI0_0)
+; LA32-JT-NEXT:    alsl.w $a3, $a3, $a4, 2
+; LA32-JT-NEXT:    ld.w $a3, $a3, 0
+; LA32-JT-NEXT:  .Ljrtb_0:
+; LA32-JT-NEXT:    jr $a3
+; LA32-JT-NEXT:  .LBB0_2: # %bb1
+; LA32-JT-NEXT:    ori $a3, $zero, 4
+; LA32-JT-NEXT:    b .LBB0_6
+; LA32-JT-NEXT:  .LBB0_3: # %bb2
+; LA32-JT-NEXT:    ori $a3, $zero, 3
+; LA32-JT-NEXT:    b .LBB0_6
+; LA32-JT-NEXT:  .LBB0_4: # %bb3
+; LA32-JT-NEXT:    ori $a3, $zero, 2
+; LA32-JT-NEXT:    b .LBB0_6
+; LA32-JT-NEXT:  .LBB0_5: # %bb4
+; LA32-JT-NEXT:    ori $a3, $zero, 1
+; LA32-JT-NEXT:  .LBB0_6: # %exit
+; LA32-JT-NEXT:    st.w $a3, $a1, 0
+; LA32-JT-NEXT:  .LBB0_7: # %exit
+; LA32-JT-NEXT:    addi.w $a3, $a0, -5
+; LA32-JT-NEXT:    bltu $a2, $a3, .LBB0_9
+; LA32-JT-NEXT:  # %bb.8: # %exit
+; LA32-JT-NEXT:    pcalau12i $a4, %pc_hi20(.LJTI0_1)
+; LA32-JT-NEXT:    addi.w $a4, $a4, %pc_lo12(.LJTI0_1)
+; LA32-JT-NEXT:    alsl.w $a3, $a3, $a4, 2
+; LA32-JT-NEXT:    ld.w $a3, $a3, 0
+; LA32-JT-NEXT:  .Ljrtb_1:
+; LA32-JT-NEXT:    jr $a3
+; LA32-JT-NEXT:  .LBB0_9: # %exit2
+; LA32-JT-NEXT:    ret
+;
+; LA64-JT-LABEL: switch_4_arms:
+; LA64-JT:       # %bb.0: # %entry
+; LA64-JT-NEXT:    addi.w $a0, $a0, 0
+; LA64-JT-NEXT:    addi.d $a3, $a0, -1
+; LA64-JT-NEXT:    ori $a2, $zero, 3
+; LA64-JT-NEXT:    bltu $a2, $a3, .LBB0_7
+; LA64-JT-NEXT:  # %bb.1: # %entry
+; LA64-JT-NEXT:    slli.d $a3, $a3, 3
+; LA64-JT-NEXT:    pcalau12i $a4, %pc_hi20(.LJTI0_0)
+; LA64-JT-NEXT:    addi.d $a4, $a4, %pc_lo12(.LJTI0_0)
+; LA64-JT-NEXT:    ldx.d $a3, $a3, $a4
+; LA64-JT-NEXT:  .Ljrtb_0:
+; LA64-JT-NEXT:    jr $a3
+; LA64-JT-NEXT:  .LBB0_2: # %bb1
+; LA64-JT-NEXT:    ori $a3, $zero, 4
+; LA64-JT-NEXT:    b .LBB0_6
+; LA64-JT-NEXT:  .LBB0_3: # %bb2
+; LA64-JT-NEXT:    ori $a3, $zero, 3
+; LA64-JT-NEXT:    b .LBB0_6
+; LA64-JT-NEXT:  .LBB0_4: # %bb3
+; LA64-JT-NEXT:    ori $a3, $zero, 2
+; LA64-JT-NEXT:    b .LBB0_6
+; LA64-JT-NEXT:  .LBB0_5: # %bb4
+; LA64-JT-NEXT:    ori $a3, $zero, 1
+; LA64-JT-NEXT:  .LBB0_6: # %exit
+; LA64-JT-NEXT:    st.w $a3, $a1, 0
+; LA64-JT-NEXT:  .LBB0_7: # %exit
+; LA64-JT-NEXT:    addi.d $a3, $a0, -5
+; LA64-JT-NEXT:    bltu $a2, $a3, .LBB0_9
+; LA64-JT-NEXT:  # %bb.8: # %exit
+; LA64-JT-NEXT:    slli.d $a3, $a3, 3
+; LA64-JT-NEXT:    pcalau12i $a4, %pc_hi20(.LJTI0_1)
+; LA64-JT-NEXT:    addi.d $a4, $a4, %pc_lo12(.LJTI0_1)
+; LA64-JT-NEXT:    ldx.d $a3, $a3, $a4
+; LA64-JT-NEXT:  .Ljrtb_1:
+; LA64-JT-NEXT:    jr $a3
+; LA64-JT-NEXT:  .LBB0_9: # %exit2
+; LA64-JT-NEXT:    ret
+entry:
+  switch i32 %in, label %exit [
+    i32 1, label %bb1
+    i32 2, label %bb2
+    i32 3, label %bb3
+    i32 4, label %bb4
+  ]
+bb1:
+  store i32 4, ptr %out
+  br label %exit
+bb2:
+  store i32 3, ptr %out
+  br label %exit
+bb3:
+  store i32 2, ptr %out
+  br label %exit
+bb4:
+  store i32 1, ptr %out
+  br label %exit
+exit:
+  switch i32 %in, label %exit2 [
+    i32 5, label %bb1
+    i32 6, label %bb2
+    i32 7, label %bb3
+    i32 8, label %bb4
+  ]
+exit2:
+  ret void
+}
+
+; UTC_ARGS: --disable
+
+; LA32-JT-LABEL: .LJTI0_0:
+; LA32-JT:       .section .discard.tablejump_annotate,"",@progbits
+; LA32-JT-NEXT:  .word .Ljrtb_0
+; LA32-JT-NEXT:  .word .LJTI0_0
+; LA32-JT-NEXT:  .word .Ljrtb_1
+; LA32-JT-NEXT:  .word .LJTI0_1
+
+; UTC_ARGS: --disable
+; LA64-JT-LABEL: .LJTI0_0:
+; LA64-JT:       .section .discard.tablejump_annotate,"",@progbits
+; LA64-JT-NEXT:  .dword .Ljrtb_0
+; LA64-JT-NEXT:  .dword .LJTI0_0
+; LA64-JT-NEXT:  .dword .Ljrtb_1
+; LA64-JT-NEXT:  .dword .LJTI0_1