[BPF] Handle traps with kfunc call __bpf_trap

llvm/lib/Target/BPF/BPF.h

@@ -22,6 +22,8 @@ class BPFTargetMachine;
 class InstructionSelector;
 class PassRegistry;
 
+static const char *BPF_TRAP = "__bpf_trap";
+
 ModulePass *createBPFCheckAndAdjustIR();
 
 FunctionPass *createBPFISelDag(BPFTargetMachine &TM);

llvm/lib/Target/BPF/BPFISelLowering.cpp

@@ -21,8 +21,10 @@
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/TargetLoweringObjectFileImpl.h"
 #include "llvm/CodeGen/ValueTypes.h"
+#include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/DiagnosticPrinter.h"
+#include "llvm/IR/Module.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/MathExtras.h"
@@ -68,6 +70,8 @@ BPFTargetLowering::BPFTargetLowering(const TargetMachine &TM,
   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::DYNAMIC_STACKALLOC, MVT::i64, Custom);
@@ -326,6 +330,8 @@ SDValue BPFTargetLowering::LowerOperation(SDValue Op, SelectionDAG &DAG) const {
   case ISD::ATOMIC_LOAD:
   case ISD::ATOMIC_STORE:
     return LowerATOMIC_LOAD_STORE(Op, DAG);
+  case ISD::TRAP:
+    return LowerTRAP(Op, DAG);
   }
 }
 
@@ -521,10 +527,12 @@ SDValue BPFTargetLowering::LowerCall(TargetLowering::CallLoweringInfo &CLI,
     Callee = DAG.getTargetGlobalAddress(G->getGlobal(), CLI.DL, PtrVT,
                                         G->getOffset(), 0);
   } else if (ExternalSymbolSDNode *E = dyn_cast<ExternalSymbolSDNode>(Callee)) {
-    Callee = DAG.getTargetExternalSymbol(E->getSymbol(), PtrVT, 0);
-    fail(CLI.DL, DAG,
-         Twine("A call to built-in function '" + StringRef(E->getSymbol()) +
-               "' is not supported."));
+    if (StringRef(E->getSymbol()) != BPF_TRAP) {
+      Callee = DAG.getTargetExternalSymbol(E->getSymbol(), PtrVT, 0);
+      fail(CLI.DL, DAG,
+           Twine("A call to built-in function '" + StringRef(E->getSymbol()) +
+                 "' is not supported."));
+    }
   }
 
   // Returns a chain & a flag for retval copy to use.
@@ -726,6 +734,52 @@ SDValue BPFTargetLowering::LowerATOMIC_LOAD_STORE(SDValue Op,
   return Op;
 }
 
+static Function *createBPFUnreachable(Module *M) {
+  if (auto *Fn = M->getFunction(BPF_TRAP))
+    return Fn;
+
+  FunctionType *FT = FunctionType::get(Type::getVoidTy(M->getContext()), false);
+  Function *NewF =
+      Function::Create(FT, GlobalValue::ExternalWeakLinkage, BPF_TRAP, M);
+  NewF->setDSOLocal(true);
+  NewF->setCallingConv(CallingConv::C);
+  NewF->setSection(".ksyms");
+
+  if (M->debug_compile_units().empty())
+    return NewF;
+
+  DIBuilder DBuilder(*M);
+  DITypeRefArray ParamTypes =
+      DBuilder.getOrCreateTypeArray({nullptr /*void return*/});
+  DISubroutineType *FuncType = DBuilder.createSubroutineType(ParamTypes);
+  DICompileUnit *CU = *M->debug_compile_units_begin();
+  DISubprogram *SP =
+      DBuilder.createFunction(CU, BPF_TRAP, BPF_TRAP, nullptr, 0, FuncType, 0,
+                              DINode::FlagZero, DISubprogram::SPFlagZero);
+  NewF->setSubprogram(SP);
+  return NewF;
+}
+
+SDValue BPFTargetLowering::LowerTRAP(SDValue Op, SelectionDAG &DAG) const {
+  MachineFunction &MF = DAG.getMachineFunction();
+  TargetLowering::CallLoweringInfo CLI(DAG);
+  SmallVector<SDValue> InVals;
+  SDNode *N = Op.getNode();
+  SDLoc DL(N);
+
+  Function *Fn = createBPFUnreachable(MF.getFunction().getParent());
+  auto PtrVT = getPointerTy(MF.getDataLayout());
+  CLI.Callee = DAG.getTargetGlobalAddress(Fn, DL, PtrVT);
+  CLI.Chain = N->getOperand(0);
+  CLI.IsTailCall = false;
+  CLI.CallConv = CallingConv::C;
+  CLI.IsVarArg = false;
+  CLI.DL = DL;
+  CLI.NoMerge = false;
+  CLI.DoesNotReturn = true;
+  return LowerCall(CLI, InVals);
+}
+
 const char *BPFTargetLowering::getTargetNodeName(unsigned Opcode) const {
   switch ((BPFISD::NodeType)Opcode) {
   case BPFISD::FIRST_NUMBER:

llvm/lib/Target/BPF/BPFISelLowering.h

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

llvm/lib/Target/BPF/BPFMIPeephole.cpp

@@ -320,6 +320,7 @@ struct BPFMIPreEmitPeephole : public MachineFunctionPass {
   bool adjustBranch();
   bool insertMissingCallerSavedSpills();
   bool removeMayGotoZero();
+  bool addExitAfterUnreachable();
 
 public:
 
@@ -336,6 +337,7 @@ struct BPFMIPreEmitPeephole : public MachineFunctionPass {
       Changed = adjustBranch() || Changed;
     Changed |= insertMissingCallerSavedSpills();
     Changed |= removeMayGotoZero();
+    Changed |= addExitAfterUnreachable();
     return Changed;
   }
 };
@@ -734,6 +736,20 @@ bool BPFMIPreEmitPeephole::removeMayGotoZero() {
   return Changed;
 }
 
+// If the last insn in a funciton is 'JAL &bpf_unreachable', let us add an
+// 'exit' insn after that insn. This will ensure no fallthrough at the last
+// insn, making kernel verification easier.
+bool BPFMIPreEmitPeephole::addExitAfterUnreachable() {
+  MachineBasicBlock &MBB = MF->back();
+  MachineInstr &MI = MBB.back();
+  if (MI.getOpcode() != BPF::JAL || !MI.getOperand(0).isGlobal() ||
+      MI.getOperand(0).getGlobal()->getName() != BPF_TRAP)
+    return false;
+
+  BuildMI(&MBB, MI.getDebugLoc(), TII->get(BPF::RET));
+  return true;
+}
+
 } // end default namespace
 
 INITIALIZE_PASS(BPFMIPreEmitPeephole, "bpf-mi-pemit-peephole",

llvm/lib/Target/BPF/BPFTargetMachine.cpp

@@ -37,6 +37,10 @@ static cl::
 opt<bool> DisableMIPeephole("disable-bpf-peephole", cl::Hidden,
                             cl::desc("Disable machine peepholes for BPF"));
 
+static cl::opt<bool>
+    DisableCheckUnreachable("bpf-disable-trap-unreachable", cl::Hidden,
+                            cl::desc("Disable Trap Unreachable for BPF"));
+
 extern "C" LLVM_EXTERNAL_VISIBILITY void LLVMInitializeBPFTarget() {
   // Register the target.
   RegisterTargetMachine<BPFTargetMachine> X(getTheBPFleTarget());
@@ -77,6 +81,11 @@ BPFTargetMachine::BPFTargetMachine(const Target &T, const Triple &TT,
                                getEffectiveCodeModel(CM, CodeModel::Small), OL),
       TLOF(std::make_unique<TargetLoweringObjectFileELF>()),
       Subtarget(TT, std::string(CPU), std::string(FS), *this) {
+  if (!DisableCheckUnreachable) {
+    this->Options.TrapUnreachable = true;
+    this->Options.NoTrapAfterNoreturn = true;
+  }
+
   initAsmInfo();
 
   BPFMCAsmInfo *MAI =

llvm/lib/Target/BPF/BTFDebug.cpp

@@ -1622,6 +1622,13 @@ void BTFDebug::endModule() {
   // Collect global types/variables except MapDef globals.
   processGlobals(false);
 
+  // In case that BPF_TRAP usage is removed during machine-level optimization,
+  // generate btf for BPF_TRAP function here.
+  for (const Function &F : *MMI->getModule()) {
+    if (F.getName() == BPF_TRAP)
+      processFuncPrototypes(&F);
+  }
+
   for (auto &DataSec : DataSecEntries)
     addType(std::move(DataSec.second));
 

llvm/test/CodeGen/BPF/BTF/builtin_trap.ll

@@ -0,0 +1,39 @@
+; RUN: llc -mtriple=bpfel -mcpu=v3 -filetype=obj -o %t1 %s
+; RUN: llvm-objcopy --dump-section='.BTF'=%t2 %t1
+; RUN: %python %p/print_btf.py %t2 | FileCheck -check-prefixes=CHECK-BTF %s
+; RUN: llc -mtriple=bpfel -mcpu=v3 < %s | FileCheck -check-prefixes=CHECK %s
+
+; BPFTargetMachine Options.NoTrapAfterNoreturn has been set to true,
+; so in below code, 'unreachable' will become a noop and
+; llvm.trap() will become 'call __bpf_trap' after selectiondag.
+define dso_local void @foo(i32 noundef %0) {
+  tail call void @llvm.trap()
+  unreachable
+}
+
+; CHECK:      .Lfunc_begin0:
+; CHECK-NEXT:   .cfi_startproc
+; CHECK-NEXT: # %bb.0:
+; CHECK-NEXT:    call __bpf_trap
+; CHECK-NEXT:    exit
+; CHECK-NEXT: .Lfunc_end0:
+
+; CHECK-BTF: [1] FUNC_PROTO '(anon)' ret_type_id=0 vlen=0
+; CHECK-BTF: [2] FUNC '__bpf_trap' type_id=1 linkage=extern
+; CHECK-BTF: [3] DATASEC '.ksyms' size=0 vlen=1
+; CHECK-BTF:  type_id=2 offset=0 size=0
+
+declare void @llvm.trap() #1
+
+attributes #1 = {noreturn}
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3, !4, !5, !6}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, emissionKind: FullDebug)
+!1 = !DIFile(filename: "test_trap.c", directory: "/some/dir")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = !{i32 1, !"wchar_size", i32 4}
+!5 = !{i32 7, !"frame-pointer", i32 2}
+!6 = !{i32 7, !"debug-info-assignment-tracking", i1 true}

llvm/test/CodeGen/BPF/BTF/unreachable.ll

@@ -0,0 +1,53 @@
+; RUN: llc -mtriple=bpfel -mcpu=v3 -filetype=obj -o %t1 %s
+; RUN: llvm-objcopy --dump-section='.BTF'=%t2 %t1
+; RUN: %python %p/print_btf.py %t2 | FileCheck -check-prefixes=CHECK-BTF %s
+; RUN: llc -mtriple=bpfel -mcpu=v3 < %s | FileCheck -check-prefixes=CHECK %s
+
+define void @foo() {
+entry:
+  tail call void @bar()
+  unreachable
+}
+
+; CHECK:      foo:
+; CHECK-NEXT: .Lfunc_begin0:
+; CHECK-NEXT:   .cfi_startproc
+; CHECK-NEXT: # %bb.0:
+; CHECK-NEXT:    call bar
+; CHECK-NEXT:    call __bpf_trap
+; CHECK-NEXT:    exit
+; CHECK-NEXT: .Lfunc_end0:
+
+define void @buz() #0 {
+entry:
+  tail call void asm sideeffect "r0 = r1; exit;", ""()
+  unreachable
+}
+
+; CHECK: buz:
+; CHECK-NEXT: .Lfunc_begin1:
+; CHECK-NEXT:   .cfi_startproc
+; CHECK-NEXT: # %bb.0:
+; CHECK-NEXT:   #APP
+; CHECK-NEXT:   r0 = r1
+; CHECK-NEXT:   exit
+; CHECK-EMPTY:
+; CHECK-NEXT:   #NO_APP
+; CHECK-NEXT: .Lfunc_end1:
+
+; CHECK-BTF: [1] FUNC_PROTO '(anon)' ret_type_id=0 vlen=0
+; CHECK-BTF: [2] FUNC '__bpf_trap' type_id=1 linkage=extern
+; CHECK-BTF: [3] DATASEC '.ksyms' size=0 vlen=1
+; CHECK-BTF:  type_id=2 offset=0 size=0
+
+declare dso_local void @bar()
+
+attributes #0 = { naked }
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!2, !3}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C11, file: !1, emissionKind: FullDebug)
+!1 = !DIFile(filename: "test.c", directory: "/some/dir")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}