[BPF] Support Jump Table (#149715)

Add jump table (switch statement and computed goto) support for BPF backend. A `gotox <reg>` insn is implemented and the `<reg>` holds the target insn where the gotox will go. For a switch statement like ``` ... 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; } ... ``` and the final binary ``` The final binary: 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 BPF.JT.0.0 7: 0f 12 00 00 00 00 00 00 r2 += r1 ... Symbol table: 4: 0000000000000000 240 OBJECT GLOBAL DEFAULT 4 BPF.JT.0.0 5: 0000000000000000 4 OBJECT GLOBAL DEFAULT 6 ret_user 6: 0000000000000000 0 NOTYPE GLOBAL DEFAULT UND bar 7: 00000000000000f0 256 OBJECT GLOBAL DEFAULT 4 BPF.JT.0.1 and [ 4] .jumptables PROGBITS 0000000000000000 0001c8 0001f0 00 0 0 1 ``` Note that for the above example, `-mllvm -bpf-min-jump-table-entries=5` should be in compilation flags as the current default bpf-min-jump-table-entries is 13. For example. ``` clang --target=bpf -mcpu=v4 -O2 -mllvm -bpf-min-jump-table-entries=5 -S -g test.c ``` For computed goto like ``` 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; } ``` The final binary: ``` 12: bf 23 20 00 00 00 00 00 r3 = (s32)r2 13: 67 03 00 00 03 00 00 00 r3 <<= 0x3 14: 18 02 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r2 = 0x0 ll 0000000000000070: R_BPF_64_64 BPF.JT.0.0 16: 0f 32 00 00 00 00 00 00 r2 += r3 17: bf 11 20 00 00 00 00 00 r1 = (s32)r1 18: 67 01 00 00 03 00 00 00 r1 <<= 0x3 19: 18 03 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r3 = 0x0 ll 0000000000000098: R_BPF_64_64 BPF.JT.0.1 21: 0f 13 00 00 00 00 00 00 r3 += r1 [ 4] .jumptables PROGBITS 0000000000000000 000160 000020 00 0 0 1 4: 0000000000000000 16 OBJECT GLOBAL DEFAULT 4 BPF.JT.0.0 5: 0000000000000010 16 OBJECT GLOBAL DEFAULT 4 BPF.JT.0.1 ``` A more complicated test with both switch-statement triggered jump table and compute gotos: ``` $ cat test3.c 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, 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; bar(); 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; } return ret; } ``` Compile with ``` clang --target=bpf -mcpu=v4 -O2 -S test3.c clang --target=bpf -mcpu=v4 -O2 -c test3.c ``` The binary: ``` /* For computed goto */ 13: bf 42 20 00 00 00 00 00 r2 = (s32)r4 14: 67 02 00 00 03 00 00 00 r2 <<= 0x3 15: 18 01 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r1 = 0x0 ll 0000000000000078: R_BPF_64_64 BPF.JT.0.1 17: 0f 21 00 00 00 00 00 00 r1 += r2 18: bf 32 20 00 00 00 00 00 r2 = (s32)r3 19: 67 02 00 00 03 00 00 00 r2 <<= 0x3 20: 18 03 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r3 = 0x0 ll 00000000000000a0: R_BPF_64_64 BPF.JT.0.2 22: 0f 23 00 00 00 00 00 00 r3 += r2 /* For switch statement */ 39: 67 01 00 00 03 00 00 00 r1 <<= 0x3 40: 18 02 00 00 00 00 00 00 00 00 00 00 00 00 00 00 r2 = 0x0 ll 0000000000000140: R_BPF_64_64 BPF.JT.0.0 42: 0f 12 00 00 00 00 00 00 r2 += r1 ``` You can see jump table symbols are all different.

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/BPFAsmPrinter.cpp

@@ -11,52 +11,35 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "BPFAsmPrinter.h"
 #include "BPF.h"
 #include "BPFInstrInfo.h"
 #include "BPFMCInstLower.h"
 #include "BTFDebug.h"
 #include "MCTargetDesc/BPFInstPrinter.h"
 #include "TargetInfo/BPFTargetInfo.h"
+#include "llvm/BinaryFormat/ELF.h"
 #include "llvm/CodeGen/AsmPrinter.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
 #include "llvm/CodeGen/MachineInstr.h"
+#include "llvm/CodeGen/MachineJumpTableInfo.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/CodeGen/TargetLowering.h"
 #include "llvm/IR/Module.h"
 #include "llvm/MC/MCAsmInfo.h"
+#include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/MC/MCSymbol.h"
+#include "llvm/MC/MCSymbolELF.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetLoweringObjectFile.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "asm-printer"
 
-namespace {
-class BPFAsmPrinter : public AsmPrinter {
-public:
- explicit BPFAsmPrinter(TargetMachine &TM,
- std::unique_ptr<MCStreamer> Streamer)
- : AsmPrinter(TM, std::move(Streamer), ID), BTF(nullptr) {}
-
- StringRef getPassName() const override { return "BPF Assembly Printer"; }
- bool doInitialization(Module &M) override;
- void printOperand(const MachineInstr *MI, int OpNum, raw_ostream &O);
- bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
- const char *ExtraCode, raw_ostream &O) override;
- bool PrintAsmMemoryOperand(const MachineInstr *MI, unsigned OpNum,
- const char *ExtraCode, raw_ostream &O) override;
-
- void emitInstruction(const MachineInstr *MI) override;
-
- static char ID;
-
-private:
- BTFDebug *BTF;
-};
-} // namespace
-
 bool BPFAsmPrinter::doInitialization(Module &M) {
  AsmPrinter::doInitialization(M);
 
@@ -69,6 +52,45 @@ bool BPFAsmPrinter::doInitialization(Module &M) {
  return false;
 }
 
+const BPFTargetMachine &BPFAsmPrinter::getBTM() const {
+ return static_cast<const BPFTargetMachine &>(TM);
+}
+
+bool BPFAsmPrinter::doFinalization(Module &M) {
+ // Remove unused globals which are previously used for jump table.
+ const BPFSubtarget *Subtarget = getBTM().getSubtargetImpl();
+ if (Subtarget->hasGotox()) {
+ std::vector<GlobalVariable *> Targets;
+ for (GlobalVariable &Global : M.globals()) {
+ if (Global.getLinkage() != GlobalValue::PrivateLinkage)
+ continue;
+ if (!Global.isConstant() || !Global.hasInitializer())
+ continue;
+
+ Constant *CV = dyn_cast<Constant>(Global.getInitializer());
+ if (!CV)
+ continue;
+ ConstantArray *CA = dyn_cast<ConstantArray>(CV);
+ if (!CA)
+ continue;
+
+ for (unsigned i = 1, e = CA->getNumOperands(); i != e; ++i) {
+ if (!dyn_cast<BlockAddress>(CA->getOperand(i)))
+ continue;
+ }
+ Targets.push_back(&Global);
+ }
+
+ for (GlobalVariable *GV : Targets) {
+ GV->replaceAllUsesWith(PoisonValue::get(GV->getType()));
+ GV->dropAllReferences();
+ GV->eraseFromParent();
+ }
+ }
+
+ return AsmPrinter::doFinalization(M);
+}
+
 void BPFAsmPrinter::printOperand(const MachineInstr *MI, int OpNum,
  raw_ostream &O) {
  const MachineOperand &MO = MI->getOperand(OpNum);
@@ -150,6 +172,50 @@ void BPFAsmPrinter::emitInstruction(const MachineInstr *MI) {
  EmitToStreamer(*OutStreamer, TmpInst);
 }
 
+MCSymbol *BPFAsmPrinter::getJTPublicSymbol(unsigned JTI) {
+ SmallString<60> Name;
+ raw_svector_ostream(Name)
+ << "BPF.JT." << MF->getFunctionNumber() << '.' << JTI;
+ MCSymbol *S = OutContext.getOrCreateSymbol(Name);
+ if (auto *ES = static_cast<MCSymbolELF *>(S)) {
+ ES->setBinding(ELF::STB_GLOBAL);
+ ES->setType(ELF::STT_OBJECT);
+ }
+ return S;
+}
+
+void BPFAsmPrinter::emitJumpTableInfo() {
+ const MachineJumpTableInfo *MJTI = MF->getJumpTableInfo();
+ if (!MJTI)
+ return;
+
+ const std::vector<MachineJumpTableEntry> &JT = MJTI->getJumpTables();
+ if (JT.empty())
+ return;
+
+ const TargetLoweringObjectFile &TLOF = getObjFileLowering();
+ const Function &F = MF->getFunction();
+ MCSection *JTS = TLOF.getSectionForJumpTable(F, TM);
+ assert(MJTI->getEntryKind() == MachineJumpTableInfo::EK_BlockAddress);
+ unsigned EntrySize = MJTI->getEntrySize(getDataLayout());
+ OutStreamer->switchSection(JTS);
+ for (unsigned JTI = 0; JTI < JT.size(); JTI++) {
+ ArrayRef<MachineBasicBlock *> JTBBs = JT[JTI].MBBs;
+ if (JTBBs.empty())
+ continue;
+
+ MCSymbol *JTStart = getJTPublicSymbol(JTI);
+ OutStreamer->emitLabel(JTStart);
+ for (const MachineBasicBlock *MBB : JTBBs) {
+ const MCExpr *LHS = MCSymbolRefExpr::create(MBB->getSymbol(), OutContext);
+ OutStreamer->emitValue(LHS, EntrySize);
+ }
+ const MCExpr *JTSize =
+ MCConstantExpr::create(JTBBs.size() * EntrySize, OutContext);
+ OutStreamer->emitELFSize(JTStart, JTSize);
+ }
+}
+
 char BPFAsmPrinter::ID = 0;
 
 INITIALIZE_PASS(BPFAsmPrinter, "bpf-asm-printer", "BPF Assembly Printer", false,

llvm/lib/Target/BPF/BPFAsmPrinter.h

@@ -0,0 +1,48 @@
+//===-- BPFFrameLowering.h - Define frame lowering for BPF -----*- C++ -*--===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_LIB_TARGET_BPF_BPFASMPRINTER_H
+#define LLVM_LIB_TARGET_BPF_BPFASMPRINTER_H
+
+#include "BPFTargetMachine.h"
+#include "BTFDebug.h"
+#include "llvm/CodeGen/AsmPrinter.h"
+
+namespace llvm {
+
+class BPFAsmPrinter : public AsmPrinter {
+public:
+ explicit BPFAsmPrinter(TargetMachine &TM,
+ std::unique_ptr<MCStreamer> Streamer)
+ : AsmPrinter(TM, std::move(Streamer), ID), BTF(nullptr), TM(TM) {}
+
+ StringRef getPassName() const override { return "BPF Assembly Printer"; }
+ bool doInitialization(Module &M) override;
+ bool doFinalization(Module &M) override;
+ void printOperand(const MachineInstr *MI, int OpNum, raw_ostream &O);
+ bool PrintAsmOperand(const MachineInstr *MI, unsigned OpNo,
+ const char *ExtraCode, raw_ostream &O) override;
+ bool PrintAsmMemoryOperand(const MachineInstr *MI, unsigned OpNum,
+ const char *ExtraCode, raw_ostream &O) override;
+
+ void emitInstruction(const MachineInstr *MI) override;
+ MCSymbol *getJTPublicSymbol(unsigned JTI);
+ virtual void emitJumpTableInfo() override;
+
+ static char ID;
+
+private:
+ BTFDebug *BTF;
+ TargetMachine &TM;
+
+ const BPFTargetMachine &getBTM() const;
+};
+
+} // namespace llvm
+
+#endif /* LLVM_LIB_TARGET_BPF_BPFASMPRINTER_H */