[BPF] Lower computed goto and generate proper jump tables

Example 2, Asm code:
        ...
    # %bb.0:                                # %entry
        r2 = .LJTI0_0 ll
        r2 = *(u64 *)(r2 + 0)
        r3 = .LJTI0_1 ll
        r3 = *(u64 *)(r3 + 0)
        if w1 == 0 goto LBB0_2
    # %bb.1:                                # %entry
        r3 = r2
    LBB0_2:                                 # %entry
        *(u64 *)(r10 - 8) = r3
        r1 = *(u64 *)(r10 - 8)
        gotox r1
    .Ltmp0:                                 # Block address taken
    LBB0_3:                                 # %l1
        w0 = 3
        goto LBB0_5
    .Ltmp1:                                 # Block address taken
    LBB0_4:                                 # %l2
        w0 = 2
    LBB0_5:                                 # %.split
        exit
        ...
        .section        .rodata,"a",@progbits
        .p2align        3, 0x0
    .LJTI0_0:
        .quad   LBB0_3
    .LJTI0_1:
        .quad   LBB0_4

Example 3, Asm Code:

        r3 = (s32)r2
        r3 <<= 3
        r2 = .LJTI0_0 ll
        r2 += r3
        r1 = (s32)r1
        r1 <<= 3
        r3 = .LJTI0_1 ll
        r3 += r1
        w0 = 0
        r1 = *(u64 *)(r3 + 0)
        gotox r1
    .Ltmp0:                                 # Block address taken
    LBB0_1:                                 # %l1
                                            # =>This Inner Loop Header: Depth=1
        w0 += 1
        w0 += 3
        r1 = *(u64 *)(r2 + 0)
        gotox r1
    .Ltmp1:                                 # Block address taken
    LBB0_2:                                 # %l2
                                            # =>This Inner Loop Header: Depth=1
        w0 += 3
        r1 = *(u64 *)(r2 + 0)
        gotox r1
    .Ltmp2:                                 # Block address taken
    LBB0_3:                                 # %l3
        w0 += 5
        goto LBB0_5
    .Ltmp3:                                 # Block address taken
    LBB0_4:                                 # %l4
    LBB0_5:                                 # %.split17
        w0 += 7
        exit
        ...
        .section        .rodata,"a",@progbits
        .p2align        3, 0x0
    .LJTI0_0:
        .quad   LBB0_3
        .quad   LBB0_4
    .LJTI0_1:
        .quad   LBB0_1
        .quad   LBB0_2
                                        # -- End function
        .type   .L__const.foo.jt1,@object       # @__const.foo.jt1
        .p2align        3, 0x0
    .L__const.foo.jt1:
        .quad   .Ltmp0
        .quad   .Ltmp1
        .size   .L__const.foo.jt1, 16

        .type   .L__const.foo.jt2,@object       # @__const.foo.jt2
        .p2align        3, 0x0
    .L__const.foo.jt2:
        .quad   .Ltmp2
        .quad   .Ltmp3
        .size   .L__const.foo.jt2, 16

Note that for both above examples, the jump table section is '.rodata'
and labels have '.L' prefix which means labels won't show up in the
symbol table. As mentioned in previous patch, we want to
  - Move jump tables to '.jumptables' section
  - Rename '.L*' labels with proper labels which are visible in symbol table.

Note that for Example 3, there are extra global functions like
  .L__const.foo.jt1 and .L__const.foo.jt2
and we are not able to remove them. But they won't show up in symbol
table either.

Author: Yonghong Song

llvm/lib/Target/BPF/BPFISelLowering.cpp

@@ -819,23 +819,12 @@ const char *BPFTargetLowering::getTargetNodeName(unsigned Opcode) const {
   return nullptr;
 }
 
-static SDValue getTargetNode(GlobalAddressSDNode *N, const SDLoc &DL, EVT Ty,
-                             SelectionDAG &DAG, unsigned Flags) {
-  return DAG.getTargetGlobalAddress(N->getGlobal(), DL, Ty, 0, Flags);
-}
-
 static SDValue getTargetNode(ConstantPoolSDNode *N, const SDLoc &DL, EVT Ty,
                              SelectionDAG &DAG, unsigned Flags) {
   return DAG.getTargetConstantPool(N->getConstVal(), Ty, N->getAlign(),
                                    N->getOffset(), Flags);
 }
 
-static SDValue getTargetNode(BlockAddressSDNode *N, const SDLoc &DL, EVT Ty,
-                             SelectionDAG &DAG, unsigned Flags) {
-  return DAG.getTargetBlockAddress(N->getBlockAddress(), Ty, N->getOffset(),
-                                   Flags);
-}
-
 static SDValue getTargetNode(JumpTableSDNode *N, const SDLoc &DL, EVT Ty,
                              SelectionDAG &DAG, unsigned Flags) {
   return DAG.getTargetJumpTable(N->getIndex(), Ty, Flags);
@@ -857,7 +846,15 @@ SDValue BPFTargetLowering::LowerGlobalAddress(SDValue Op,
   if (N->getOffset() != 0)
     report_fatal_error("invalid offset for global address: " +
                        Twine(N->getOffset()));
-  return getAddr(N, DAG);
+
+  const GlobalValue *GVal = N->getGlobal();
+  SDLoc DL(Op);
+
+  // Wrap it in a TargetGlobalAddress
+  SDValue Addr = DAG.getTargetGlobalAddress(GVal, DL, MVT::i64);
+
+  // Emit pseudo instruction
+  return SDValue(DAG.getMachineNode(BPF::LDIMM64, DL, MVT::i64, Addr), 0);
 }
 
 SDValue BPFTargetLowering::LowerConstantPool(SDValue Op,
@@ -869,8 +866,14 @@ SDValue BPFTargetLowering::LowerConstantPool(SDValue Op,
 
 SDValue BPFTargetLowering::LowerBlockAddress(SDValue Op,
                                              SelectionDAG &DAG) const {
-  BlockAddressSDNode *N = cast<BlockAddressSDNode>(Op);
-  return getAddr(N, DAG);
+  const BlockAddress *BA = cast<BlockAddressSDNode>(Op)->getBlockAddress();
+  SDLoc DL(Op);
+
+  // Wrap it in a TargetBlockAddress
+  SDValue Addr = DAG.getTargetBlockAddress(BA, MVT::i64);
+
+  // Emit pseudo instruction
+  return SDValue(DAG.getMachineNode(BPF::LDIMM64, DL, MVT::i64, Addr), 0);
 }
 
 unsigned
@@ -936,6 +939,86 @@ BPFTargetLowering::EmitInstrWithCustomInserterMemcpy(MachineInstr &MI,
   return BB;
 }
 
+MachineBasicBlock *BPFTargetLowering::EmitInstrWithCustomInserterLDimm64(
+    MachineInstr &MI, MachineBasicBlock *BB) const {
+  MachineFunction *MF = BB->getParent();
+  const BPFInstrInfo *TII = MF->getSubtarget<BPFSubtarget>().getInstrInfo();
+  const TargetRegisterClass *RC = getRegClassFor(MVT::i64);
+  MachineRegisterInfo &RegInfo = MF->getRegInfo();
+  DebugLoc DL = MI.getDebugLoc();
+
+  // Build address taken map for Global Varaibles and BlockAddresses
+  DenseMap<const BasicBlock *, MachineBasicBlock *> AddressTakenBBs;
+  for (MachineBasicBlock &MBB : *MF) {
+    if (const BasicBlock *BB = MBB.getBasicBlock())
+      if (BB->hasAddressTaken())
+        AddressTakenBBs[BB] = &MBB;
+  }
+
+  MachineOperand &MO = MI.getOperand(1);
+  assert(MO.isBlockAddress() || MO.isGlobal());
+
+  MCRegister ResultReg = MI.getOperand(0).getReg();
+  Register TmpReg = RegInfo.createVirtualRegister(RC);
+
+  std::vector<MachineBasicBlock *> Targets;
+  unsigned JTI;
+
+  if (MO.isBlockAddress()) {
+    auto *BA = MO.getBlockAddress();
+    MachineBasicBlock *TgtMBB = AddressTakenBBs[BA->getBasicBlock()];
+    assert(TgtMBB);
+
+    Targets.push_back(TgtMBB);
+    JTI = MF->getOrCreateJumpTableInfo(getJumpTableEncoding())
+              ->createJumpTableIndex(Targets);
+
+    BuildMI(*BB, MI, DL, TII->get(BPF::LD_imm64), TmpReg)
+        .addJumpTableIndex(JTI);
+    BuildMI(*BB, MI, DL, TII->get(BPF::LDD), ResultReg)
+        .addReg(TmpReg)
+        .addImm(0);
+    MI.eraseFromParent();
+    return BB;
+  }
+
+  // Helper: emit LD_imm64 with operand GlobalAddress or JumpTable
+  auto emitLDImm64 = [&](const GlobalValue *GV = nullptr, unsigned JTI = -1) {
+    auto MIB = BuildMI(*BB, MI, DL, TII->get(BPF::LD_imm64), ResultReg);
+    if (GV)
+      MIB.addGlobalAddress(GV);
+    else
+      MIB.addJumpTableIndex(JTI);
+    MI.eraseFromParent();
+    return BB;
+  };
+
+  // Must be a global at this point
+  const GlobalValue *GVal = MO.getGlobal();
+  const auto *GV = dyn_cast<GlobalVariable>(GVal);
+
+  if (!GV || GV->getLinkage() != GlobalValue::PrivateLinkage ||
+      !GV->isConstant() || !GV->hasInitializer())
+    return emitLDImm64(GVal);
+
+  const auto *CA = dyn_cast<ConstantArray>(GV->getInitializer());
+  if (!CA)
+    return emitLDImm64(GVal);
+
+  for (const Use &Op : CA->operands()) {
+    if (!isa<BlockAddress>(Op))
+      return emitLDImm64(GVal);
+    auto *BA = cast<BlockAddress>(Op);
+    MachineBasicBlock *TgtMBB = AddressTakenBBs[BA->getBasicBlock()];
+    assert(TgtMBB);
+    Targets.push_back(TgtMBB);
+  }
+
+  JTI = MF->getOrCreateJumpTableInfo(getJumpTableEncoding())
+            ->createJumpTableIndex(Targets);
+  return emitLDImm64(nullptr, JTI);
+}
+
 MachineBasicBlock *
 BPFTargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
                                                MachineBasicBlock *BB) const {
@@ -948,20 +1031,24 @@ BPFTargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
                        Opc == BPF::Select_32_64);
 
   bool isMemcpyOp = Opc == BPF::MEMCPY;
+  bool isLDimm64Op = Opc == BPF::LDIMM64;
 
 #ifndef NDEBUG
   bool isSelectRIOp = (Opc == BPF::Select_Ri ||
                        Opc == BPF::Select_Ri_64_32 ||
                        Opc == BPF::Select_Ri_32 ||
                        Opc == BPF::Select_Ri_32_64);
 
-  if (!(isSelectRROp || isSelectRIOp || isMemcpyOp))
+  if (!(isSelectRROp || isSelectRIOp || isMemcpyOp || isLDimm64Op))
     report_fatal_error("unhandled instruction type: " + Twine(Opc));
 #endif
 
   if (isMemcpyOp)
     return EmitInstrWithCustomInserterMemcpy(MI, BB);
 
+  if (isLDimm64Op)
+    return EmitInstrWithCustomInserterLDimm64(MI, BB);
+
   bool is32BitCmp = (Opc == BPF::Select_32 ||
                      Opc == BPF::Select_32_64 ||
                      Opc == BPF::Select_Ri_32 ||

llvm/lib/Target/BPF/BPFISelLowering.h

@@ -167,7 +167,9 @@ class BPFTargetLowering : public TargetLowering {
   MachineBasicBlock * EmitInstrWithCustomInserterMemcpy(MachineInstr &MI,
                                                         MachineBasicBlock *BB)
                                                         const;
-
+  MachineBasicBlock *
+  EmitInstrWithCustomInserterLDimm64(MachineInstr &MI,
+                                     MachineBasicBlock *BB) const;
 };
 }
 

llvm/lib/Target/BPF/BPFInstrInfo.td

@@ -865,9 +865,7 @@ let usesCustomInserter = 1, isCodeGenOnly = 1 in {
 }
 
 // load 64-bit global addr into register
-def : Pat<(BPFWrapper tglobaladdr:$in), (LD_imm64 tglobaladdr:$in)>;
 def : Pat<(BPFWrapper tconstpool:$in), (LD_imm64 tconstpool:$in)>;
-def : Pat<(BPFWrapper tblockaddress:$in), (LD_imm64 tblockaddress:$in)>;
 def : Pat<(BPFWrapper tjumptable:$in), (LD_imm64 tjumptable:$in)>;
 
 // 0xffffFFFF doesn't fit into simm32, optimize common case
@@ -1388,3 +1386,8 @@ let usesCustomInserter = 1, isCodeGenOnly = 1 in {
       "#memcpy dst: $dst, src: $src, len: $len, align: $align",
       [(BPFmemcpy GPR:$dst, GPR:$src, imm:$len, imm:$align)]>;
 }
+
+// For GlobalValue and BlockAddress.
+let usesCustomInserter = 1, isCodeGenOnly = 1 in {
+  def LDIMM64 : Pseudo<(outs GPR:$dst), (ins i64imm:$addr), "", []>;
+}

llvm/lib/Target/BPF/BPFMCInstLower.cpp

@@ -80,10 +80,6 @@ void BPFMCInstLower::Lower(const MachineInstr *MI, MCInst &OutMI) const {
     case MachineOperand::MO_JumpTableIndex:
       MCOp = LowerSymbolOperand(MO, Printer.GetJTISymbol(MO.getIndex()));
       break;
-    case MachineOperand::MO_BlockAddress:
-      MCOp = LowerSymbolOperand(
-          MO, Printer.GetBlockAddressSymbol(MO.getBlockAddress()));
-      break;
     }
 
     OutMI.addOperand(MCOp);