[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. Let us use cpu v4 for now.
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 -mcpu=v4 -O2 -S test.c
The assembly code:
        ...
    # %bb.1:                                # %entry
        w1 = 18
        r2 <<= 2
        r3 = .LJTI0_0 ll
        r3 += r2
        r2 = *(s32 *)(r3 + 0)
        gotox r2
    LBB0_2:                                 # %sw.bb1
        w1 = 6
        goto LBB0_8
        ...
    # %bb.9:                                # %sw.epilog
        w1 = 8
        r2 <<= 2
        r3 = .LJTI0_1 ll
        r3 += r2
        r2 = *(s32 *)(r3 + 0)
        gotox r2
    LBB0_10:                                # %sw.bb8
        w1 = 5
        goto LBB0_14
        ...
        .section        .rodata,"a",@progbits
        .p2align        2, 0x0
    .LJTI0_0:
        .long   LBB0_8-.LJTI0_0
        .long   LBB0_7-.LJTI0_0
        ...
        .long   LBB0_6-.LJTI0_0
    .LJTI0_1:
        .long   LBB0_14-.LJTI0_1
        .long   LBB0_11-.LJTI0_1
        ...

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 other symbols.

Run: llvm-objdump -Sr test.o
       ...
       5:       67 02 00 00 02 00 00 00 r2 <<= 0x2
       6:       18 03 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r3 = 0x0 ll
                0000000000000030:  R_BPF_64_64  .rodata
       8:       0f 23 00 00 00 00 00 00 r3 += r2
       ...
      29:       67 02 00 00 02 00 00 00 r2 <<= 0x2
      30:       18 03 00 00 78 00 00 00 00 00 00 00 00 00 00 00 r3 = 0x78 ll
                00000000000000f0:  R_BPF_64_64  .rodata
      32:       0f 23 00 00 00 00 00 00 r3 += r2

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 -mcpu=v4 -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 -mcpu=v4 -O2 -S test2.c
Asm code:
        ...
        r3 = (s32)r2
        r3 <<= 3
        r2 = .L__const.foo.jt2 ll
        r2 += r3
        r1 = (s32)r1
        r1 <<= 3
        r3 = .L__const.foo.jt1 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
        ...
        .type   .L__const.foo.jt1,@object       # @__const.foo.jt1
        .section        .rodata,"a",@progbits
        .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

Similar to switch statement case, for the binary, the symbols
.L__const.foo.jt* 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

@@ -18,6 +18,7 @@
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/ValueTypes.h"
@@ -38,6 +39,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 +72,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 +165,7 @@ BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
 
   setBooleanContents(ZeroOrOneBooleanContent);
   setMaxAtomicSizeInBitsSupported(64);
+  setMinimumJumpTableEntries(BPFMinimumJumpTableEntries);
 
   // Function alignments
   setMinFunctionAlignment(Align(8));
@@ -246,6 +253,10 @@ bool BPFTargetLowering::isZExtFree(SDValue Val, EVT VT2) const {
   return TargetLoweringBase::isZExtFree(Val, VT2);
 }
 
+unsigned BPFTargetLowering::getJumpTableEncoding() const {
+  return MachineJumpTableInfo::EK_LabelDifference32;
+}
+
 BPFTargetLowering::ConstraintType
 BPFTargetLowering::getConstraintType(StringRef Constraint) const {
   if (Constraint.size() == 1) {
@@ -316,10 +327,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 +795,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 +831,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 +868,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

@@ -66,6 +66,8 @@ class BPFTargetLowering : public TargetLowering {
 
   MVT getScalarShiftAmountTy(const DataLayout &, EVT) const override;
 
+  unsigned getJumpTableEncoding() const override;
+
 private:
   // Control Instruction Selection Features
   bool HasAlu32;
@@ -81,6 +83,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,11 @@ bool BPFInstrInfo::analyzeBranch(MachineBasicBlock &MBB,
     if (!isUnpredicatedTerminator(*I))
       break;
 
+    // From base method doc: ... returning true if it cannot be understood ...
+    // Indirect branch has multiple destinations and no true/false concepts.
+    if (I->isIndirectBranch())
+      return true;
+
     // A terminator that isn't a branch can't easily be handled
     // by this analysis.
     if (!I->isBranch())
@@ -259,3 +264,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);