[BPF] Add jump table support with switch statements and computed goto

NOTE 1: We probably need cpu v5 or other flags to enable this feature.
We can add it later when necessary.
NOTE 2: An option -bpf-min-jump-table-entries is implemented to control the minimum
number of entries to use a jump table on BPF. The default value 5 and this is
to make it easy to test. Eventually we will increase min jump table entries to be 13.

This patch adds jump table support. A new insn 'gotox <reg>' is
added to allow goto through a register. The register represents
the address in the current section.

Example 1 (switch statement):
=============================

Code:
struct simple_ctx {
        int x;
        int y;
        int z;
};

int ret_user, ret_user2;
void bar(void);
int foo(struct simple_ctx *ctx, struct simple_ctx *ctx2)
{
        switch (ctx->x) {
        case 1: ret_user = 18; break;
        case 20: ret_user = 6; break;
        case 16: ret_user = 9; break;
        case 6: ret_user = 16; break;
        case 8: ret_user = 14; break;
        case 30: ret_user = 2; break;
        default: ret_user = 1; break;
        }

        bar();

        switch (ctx2->x) {
        case 0: ret_user2 = 8; break;
        case 31: ret_user2 = 5; break;
        case 13: ret_user2 = 8; break;
        case 1: ret_user2 = 3; break;
        case 11: ret_user2 = 4; break;
        default: ret_user2 = 29; break;
        }

        return 0;
}

Run: clang --target=bpf -O2 -S test.c
The assembly code:
  ...
  # %bb.1:                                # %entry
        r1 <<= 3
        r2 = .LJTI0_0 ll
        r2 += r1
        r1 = *(u64 *)(r2 + 0)
        gotox r1
        goto LBB0_8
  LBB0_2:
        w1 = 18
        goto LBB0_9
  ...
  # %bb.10:                               # %sw.epilog
        r1 <<= 3
        r2 = .LJTI0_1 ll
        r2 += r1
        r1 = *(u64 *)(r2 + 0)
        gotox r1
        goto LBB0_15
  ...
        .section        .rodata,"a",@progbits
        .p2align        3, 0x0
  .LJTI0_0:
        .quad   LBB0_2
        .quad   LBB0_8
        .quad   LBB0_8
        .quad   LBB0_8
        .quad   LBB0_7
        ...
  .LJTI0_1:
        .quad   LBB0_11
        .quad   LBB0_13

Although we do have labels .LJTI0_0 and .LJTI0_1, but since they have
prefix '.L' so they won't appear in the .o file like

Run: llvm-objdump -Sr test.o
  ...
       4:       67 01 00 00 03 00 00 00 r1 <<= 0x3
       5:       18 02 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r2 = 0x0 ll
                0000000000000028:  R_BPF_64_64  .rodata
       7:       0f 12 00 00 00 00 00 00 r2 += r1
  ...
      30:       67 01 00 00 03 00 00 00 r1 <<= 0x3
      31:       18 02 00 00 f0 00 00 00 00 00 00 00 00 00 00 00 r2 = 0xf0 ll
                00000000000000f8:  R_BPF_64_64  .rodata
The size of jump table is not obvious. The libbpf needs to check all relocations
against .rodata section in order to get precise size in order to construct bpf
maps.

Example 2 (Simple computed goto):
=================================

Code:
     int bar(int a) {
        __label__ l1, l2;
        void * volatile tgt;
        int ret = 0;
        if (a)
          tgt = &&l1; // synthetic jump table generated here
        else
          tgt = &&l2; // another synthetic jump table
        goto *tgt;
    l1: ret += 1;
    l2: ret += 2;
        return ret;
      }

Compile: clang --target=bpf -O2 -c test1.c
Objdump: llvm-objdump -Sr test1.o

       0:       18 02 00 00 50 00 00 00 00 00 00 00 00 00 00 00 r2 = 0x50 ll
                0000000000000000:  R_BPF_64_64  .text
       2:       16 01 02 00 00 00 00 00 if w1 == 0x0 goto +0x2 <bar+0x28>
       3:       18 02 00 00 40 00 00 00 00 00 00 00 00 00 00 00 r2 = 0x40 ll
                0000000000000018:  R_BPF_64_64  .text
       5:       7b 2a f8 ff 00 00 00 00 *(u64 *)(r10 - 0x8) = r2
       6:       79 a1 f8 ff 00 00 00 00 r1 = *(u64 *)(r10 - 0x8)
       7:       0d 01 00 00 00 00 00 00 gotox r1
       8:       b4 00 00 00 03 00 00 00 w0 = 0x3
       9:       05 00 01 00 00 00 00 00 goto +0x1 <bar+0x58>
      10:       b4 00 00 00 02 00 00 00 w0 = 0x2
      11:       95 00 00 00 00 00 00 00 exit

For this case, no need for jump table. Verifier is able to detect precise
gotox addresses for different branches.

Example 3 (More complicated computed goto):
===========================================

Code:
  int foo(int a, int b) {
    __label__ l1, l2, l3, l4;
    void *jt1[] = {[0]=&&l1, [1]=&&l2};
    void *jt2[] = {[0]=&&l3, [1]=&&l4};
    int ret = 0;

    goto *jt1[a % 2];
    l1: ret += 1;
    l2: ret += 3;
    goto *jt2[b % 2];
    l3: ret += 5;
    l4: ret += 7;
    return ret;
  }

Compile: clang --target=bpf -O2 -S test2.c
Asm code:
  # %bb.1:                                # %entry
        r1 <<= 3
        r2 = .LJTI0_0 ll
        r2 += r1
        r1 = *(u64 *)(r2 + 0)
        gotox r1
        goto LBB0_8
  LBB0_2:
        w1 = 18
        goto LBB0_9
  ...
  # %bb.10:                               # %sw.epilog
        r1 <<= 3
        r2 = .LJTI0_1 ll
        r2 += r1
        r1 = *(u64 *)(r2 + 0)
        gotox r1
        goto LBB0_15
  LBB0_11:
        w1 = 8
        goto LBB0_16
  ...
        .section        .rodata,"a",@progbits
        .p2align        3, 0x0
  .LJTI0_0:
        .quad   LBB0_2
        .quad   LBB0_8
        .quad   LBB0_8
  ...
        .quad   LBB0_7
  .LJTI0_1:
        .quad   LBB0_11
        .quad   LBB0_13
        .quad   LBB0_15
  ...

Similar to switch statement case, for the binary, the symbols .LJTI0_*
will not show up in the symbol table and jump table will be in .rodata
section.

With more libbpf work (dealing with .rodata sections etc.),
everything should work fine. But we could do better by
  - Replacing symbols like .L<...> with symbols appearing in
    symbol table.
  - Add jump tables to .jumptables section instead of .rodata section.
This should make things easier for libbpf and for users.

Next two patches try to implement the above.

Author: Yonghong Song

llvm/lib/Target/BPF/AsmParser/BPFAsmParser.cpp

@@ -234,6 +234,7 @@ struct BPFOperand : public MCParsedAsmOperand {
         .Case("callx", true)
         .Case("goto", true)
         .Case("gotol", true)
+        .Case("gotox", true)
         .Case("may_goto", true)
         .Case("*", true)
         .Case("exit", true)

llvm/lib/Target/BPF/BPFISelLowering.cpp

@@ -38,6 +38,10 @@ static cl::opt<bool> BPFExpandMemcpyInOrder("bpf-expand-memcpy-in-order",
   cl::Hidden, cl::init(false),
   cl::desc("Expand memcpy into load/store pairs in order"));
 
+static cl::opt<unsigned> BPFMinimumJumpTableEntries(
+    "bpf-min-jump-table-entries", cl::init(5), cl::Hidden,
+    cl::desc("Set minimum number of entries to use a jump table on BPF"));
+
 static void fail(const SDLoc &DL, SelectionDAG &DAG, const Twine &Msg,
                  SDValue Val = {}) {
   std::string Str;
@@ -67,12 +71,13 @@ BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
 
   setOperationAction(ISD::BR_CC, MVT::i64, Custom);
   setOperationAction(ISD::BR_JT, MVT::Other, Expand);
-  setOperationAction(ISD::BRIND, MVT::Other, Expand);
   setOperationAction(ISD::BRCOND, MVT::Other, Expand);
 
   setOperationAction(ISD::TRAP, MVT::Other, Custom);
 
-  setOperationAction({ISD::GlobalAddress, ISD::ConstantPool}, MVT::i64, Custom);
+  setOperationAction({ISD::GlobalAddress, ISD::ConstantPool, ISD::JumpTable,
+                      ISD::BlockAddress},
+                     MVT::i64, Custom);
 
   setOperationAction(ISD::DYNAMIC_STACKALLOC, MVT::i64, Custom);
   setOperationAction(ISD::STACKSAVE, MVT::Other, Expand);
@@ -159,6 +164,7 @@ BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
 
   setBooleanContents(ZeroOrOneBooleanContent);
   setMaxAtomicSizeInBitsSupported(64);
+  setMinimumJumpTableEntries(BPFMinimumJumpTableEntries);
 
   // Function alignments
   setMinFunctionAlignment(Align(8));
@@ -316,10 +322,14 @@ SDValue BPFTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
     report_fatal_error("unimplemented opcode: " + Twine(Op.getOpcode()));
   case ISD::BR_CC:
     return LowerBR_CC(Op, DAG);
+  case ISD::JumpTable:
+    return LowerJumpTable(Op, DAG);
   case ISD::GlobalAddress:
     return LowerGlobalAddress(Op, DAG);
   case ISD::ConstantPool:
     return LowerConstantPool(Op, DAG);
+  case ISD::BlockAddress:
+    return LowerBlockAddress(Op, DAG);
   case ISD::SELECT_CC:
     return LowerSELECT_CC(Op, DAG);
   case ISD::SDIV:
@@ -780,6 +790,11 @@ SDValue BPFTargetLowering::LowerTRAP(SDValue Op, SelectionDAG &DAG) const {
   return LowerCall(CLI, InVals);
 }
 
+SDValue BPFTargetLowering::LowerJumpTable(SDValue Op, SelectionDAG &DAG) const {
+  JumpTableSDNode *N = cast<JumpTableSDNode>(Op);
+  return getAddr(N, DAG);
+}
+
 const char *BPFTargetLowering::getTargetNodeName(unsigned Opcode) const {
   switch ((BPFISD::NodeType)Opcode) {
   case BPFISD::FIRST_NUMBER:
@@ -811,6 +826,17 @@ static SDValue getTargetNode(ConstantPoolSDNode *N, const SDLoc &DL, EVT Ty,
                                    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);
+}
+
 template <class NodeTy>
 SDValue BPFTargetLowering::getAddr(NodeTy *N, SelectionDAG &DAG,
                                    unsigned Flags) const {
@@ -837,6 +863,12 @@ SDValue BPFTargetLowering::LowerConstantPool(SDValue Op,
   return getAddr(N, DAG);
 }
 
+SDValue BPFTargetLowering::LowerBlockAddress(SDValue Op,
+                                             SelectionDAG &DAG) const {
+  BlockAddressSDNode *N = cast<BlockAddressSDNode>(Op);
+  return getAddr(N, DAG);
+}
+
 unsigned
 BPFTargetLowering::EmitSubregExt(MachineInstr &MI, MachineBasicBlock *BB,
                                  unsigned Reg, bool isSigned) const {

llvm/lib/Target/BPF/BPFISelLowering.h

@@ -81,6 +81,8 @@ class BPFTargetLowering : public TargetLowering {
   SDValue LowerConstantPool(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerGlobalAddress(SDValue Op, SelectionDAG &DAG) const;
   SDValue LowerTRAP(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerBlockAddress(SDValue Op, SelectionDAG &DAG) const;
+  SDValue LowerJumpTable(SDValue Op, SelectionDAG &DAG) const;
 
   template <class NodeTy>
   SDValue getAddr(NodeTy *N, SelectionDAG &DAG, unsigned Flags = 0) const;

llvm/lib/Target/BPF/BPFInstrInfo.cpp

@@ -181,6 +181,10 @@ bool BPFInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
     if (!isUnpredicatedTerminator(*I))
       break;
 
+    // If a JX insn, we're done.
+    if (I->getOpcode() == BPF::JX)
+      break;
+
     // A terminator that isn't a branch can't easily be handled
     // by this analysis.
     if (!I->isBranch())
@@ -259,3 +263,40 @@ unsigned BPFInstrInfo::removeBranch(MachineBasicBlock &MBB,
 
   return Count;
 }
+
+int BPFInstrInfo::getJumpTableIndex(const MachineInstr &MI) const {
+  // The pattern looks like:
+  // %0 = LD_imm64 %jump-table.0   ; load jump-table address
+  // %1 = ADD_rr %0, $another_reg  ; address + offset
+  // %2 = LDD %1, 0                ; load the actual label
+  // JX %2
+  const MachineFunction &MF = *MI.getParent()->getParent();
+  const MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  Register Reg = MI.getOperand(0).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *Ldd = MRI.getUniqueVRegDef(Reg);
+  if (Ldd == nullptr || Ldd->getOpcode() != BPF::LDD)
+    return -1;
+
+  Reg = Ldd->getOperand(1).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *Add = MRI.getUniqueVRegDef(Reg);
+  if (Add == nullptr || Add->getOpcode() != BPF::ADD_rr)
+    return -1;
+
+  Reg = Add->getOperand(1).getReg();
+  if (!Reg.isVirtual())
+    return -1;
+  MachineInstr *LDimm64 = MRI.getUniqueVRegDef(Reg);
+  if (LDimm64 == nullptr || LDimm64->getOpcode() != BPF::LD_imm64)
+    return -1;
+
+  const MachineOperand &MO = LDimm64->getOperand(1);
+  if (!MO.isJTI())
+    return -1;
+
+  return MO.getIndex();
+}

llvm/lib/Target/BPF/BPFInstrInfo.h

@@ -58,6 +58,9 @@ class BPFInstrInfo : public BPFGenInstrInfo {
                         MachineBasicBlock *FBB, ArrayRef<MachineOperand> Cond,
                         const DebugLoc &DL,
                         int *BytesAdded = nullptr) const override;
+
+  int getJumpTableIndex(const MachineInstr &MI) const override;
+
 private:
   void expandMEMCPY(MachineBasicBlock::iterator) const;
 

llvm/lib/Target/BPF/BPFInstrInfo.td

@@ -216,6 +216,18 @@ class JMP_RI<BPFJumpOp Opc, string OpcodeStr, PatLeaf Cond>
   let BPFClass = BPF_JMP;
 }
 
+class JMP_IND<BPFJumpOp Opc, string OpcodeStr, list<dag> Pattern>
+    : TYPE_ALU_JMP<Opc.Value, BPF_X.Value,
+                   (outs),
+                   (ins GPR:$dst),
+                   !strconcat(OpcodeStr, " $dst"),
+                   Pattern> {
+  bits<4> dst;
+
+  let Inst{51-48} = dst;
+  let BPFClass = BPF_JMP;
+}
+
 class JMP_JCOND<BPFJumpOp Opc, string OpcodeStr, list<dag> Pattern>
     : TYPE_ALU_JMP<Opc.Value, BPF_K.Value,
                    (outs),
@@ -281,6 +293,10 @@ defm JSLT : J<BPF_JSLT, "s<", BPF_CC_LT, BPF_CC_LT_32>;
 defm JSLE : J<BPF_JSLE, "s<=", BPF_CC_LE, BPF_CC_LE_32>;
 defm JSET : J<BPF_JSET, "&", NoCond, NoCond>;
 def JCOND : JMP_JCOND<BPF_JCOND, "may_goto", []>;
+
+let isIndirectBranch = 1 in {
+  def JX : JMP_IND<BPF_JA, "gotox", [(brind i64:$dst)]>;
+}
 }
 
 // ALU instructions
@@ -851,6 +867,8 @@ 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
 def : Pat<(i64 (and (i64 GPR:$src), 0xffffFFFF)),

llvm/lib/Target/BPF/BPFMCInstLower.cpp

@@ -77,6 +77,13 @@ void BPFMCInstLower::Lower(const MachineInstr *MI, MCInst &OutMI) const {
     case MachineOperand::MO_ConstantPoolIndex:
       MCOp = LowerSymbolOperand(MO, Printer.GetCPISymbol(MO.getIndex()));
       break;
+    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);