[llvm] [BPF] Support Jump Table (PR #149715)

[via llvm-commits]

================
@@ -750,6 +753,54 @@ bool BPFMIPreEmitPeephole::addExitAfterUnreachable() {
   return true;
 }
 
+bool BPFMIPreEmitPeephole::convertBAToConstantArray() {
----------------
eddyz87 wrote:

Nit: below works for me and avoids introducing temporary rodata variable.
I wonder, if the same logic can be moved to `BPFTargetLowering::LowerBlockAddress`.

```diff
--- a/llvm/lib/Target/BPF/BPFMIPeephole.cpp
+++ b/llvm/lib/Target/BPF/BPFMIPeephole.cpp
@@ -29,6 +29,7 @@
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/Support/Debug.h"
@@ -754,9 +755,15 @@ bool BPFMIPreEmitPeephole::addExitAfterUnreachable() {
 }
 
 bool BPFMIPreEmitPeephole::convertBAToConstantArray() {
+  DenseMap<const BasicBlock *, MachineBasicBlock *> AddressTakenBBs;
   bool Changed = false;
 
-  unsigned gv_idx = 0;
+  for (MachineBasicBlock &MBB : *MF) {
+    if (!MBB.getBasicBlock()->hasAddressTaken())
+      continue;
+    AddressTakenBBs[MBB.getBasicBlock()] = &MBB;
+  }
+
   for (MachineBasicBlock &MBB : *MF) {
     for (MachineInstr &MI : make_early_inc_range(MBB)) {
       if (MI.getOpcode() != BPF::LD_imm64)
@@ -772,22 +779,18 @@ bool BPFMIPreEmitPeephole::convertBAToConstantArray() {
       // BloackAddress, and then generate a global variable based on this array.
       // This will allow generating a jump table later.
       const BlockAddress *OldBA = MO.getBlockAddress();
-      BlockAddress *NewBA =
-          BlockAddress::get(OldBA->getFunction(), OldBA->getBasicBlock());
-      ArrayType *ArrTy = ArrayType::get(NewBA->getType(), 1);
-      SmallVector<Constant *, 4> Elems;
-      Elems.push_back((Constant *)NewBA);
-      Constant *Init = ConstantArray::get(ArrTy, Elems);
-      GlobalVariable *GV = new GlobalVariable(
-          *MF->getFunction().getParent(), ArrTy, true,
-          GlobalValue::PrivateLinkage, Init, "ba2gv_" + std::to_string(gv_idx));
-      gv_idx++;
+      MachineBasicBlock *TgtMBB = AddressTakenBBs[OldBA->getBasicBlock()];
+      std::vector<MachineBasicBlock *> Targets;
+      Targets.push_back(TgtMBB);
+      unsigned JTI =
+        MF->getOrCreateJumpTableInfo(MachineJumpTableInfo::EK_LabelDifference64)
+          ->createJumpTableIndex(Targets);
 
       // From 'reg = LD_imm64 <blockaddress>' to 'reg = LD_imm64 ba2gv_<idx>;
       // reg = *(u64 *)(reg + 0)'.
       BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(BPF::LD_imm64))
           .addReg(ResultReg)
-          .addGlobalAddress(GV);
+          .addJumpTableIndex(JTI);
       BuildMI(MBB, MI, MI.getDebugLoc(), TII->get(BPF::LDD))
           .addReg(ResultReg)
           .addReg(ResultReg)
```

https://github.com/llvm/llvm-project/pull/149715