Codegen: Switch to LLVM Orc-based compiler

Removes dependency on glibc for loading dynamic libraries at runtime
e.g. as in EngineBuilder::create()

src/CMakeLists.txt

@@ -16,7 +16,7 @@ add_executable(bpftrace
 
 target_link_libraries(bpftrace arch ast parser)
 
-llvm_map_components_to_libnames(llvm_libs bpfcodegen ipo irreader mcjit)
+llvm_map_components_to_libnames(llvm_libs bpfcodegen ipo irreader mcjit orcjit)
 target_link_libraries(bpftrace ${llvm_libs})
 
 add_dependencies(bpftrace bcc-build)

src/ast/codegen_llvm.cpp

+#include "bpforc.h"
 #include "codegen_llvm.h"
 #include "ast.h"
 #include "parser.tab.hh"
 #include "arch/arch.h"
 
-#include <llvm/ExecutionEngine/SectionMemoryManager.h>
 #include <llvm/Support/raw_os_ostream.h>
 #include <llvm/Support/TargetRegistry.h>
 #include <llvm/IR/LegacyPassManager.h>
@@ -562,29 +562,6 @@ Value *CodegenLLVM::createLogicalOr(Binop &binop)
   return b_.CreateLoad(result);
 }
 
-class BPFtraceMemoryManager : public SectionMemoryManager
-{
-public:
-  explicit BPFtraceMemoryManager(std::map<std::string, std::tuple<uint8_t *, uintptr_t>> &sections)
-    : sections_(sections) { }
-  uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment, unsigned SectionID, StringRef SectionName) override
-  {
-    uint8_t *addr = SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID, SectionName);
-    sections_[SectionName.str()] = std::make_tuple(addr, Size);
-    return addr;
-  }
-
-  uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment, unsigned SectionID, StringRef SectionName, bool isReadOnly) override
-  {
-    uint8_t *addr = SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID, SectionName, isReadOnly);
-    sections_[SectionName.str()] = std::make_tuple(addr, Size);
-    return addr;
-  }
-
-private:
-  std::map<std::string, std::tuple<uint8_t *, uintptr_t>> &sections_;
-};
-
 void CodegenLLVM::createLog2Function()
 {
   // log2(int n)
@@ -666,7 +643,7 @@ void CodegenLLVM::createStrcmpFunction()
   b_.CreateRet(b_.getInt1(0));
 }
 
-int CodegenLLVM::compile(bool debug, std::ostream &out)
+std::unique_ptr<BpfOrc> CodegenLLVM::compile(bool debug, std::ostream &out)
 {
   createLog2Function();
   createStrcmpFunction();
@@ -682,11 +659,8 @@ int CodegenLLVM::compile(bool debug, std::ostream &out)
 
   std::string error;
   const Target *target = TargetRegistry::lookupTarget(targetTriple, error);
-  if (!target) {
-    std::cerr << "Could not create LLVM target" << std::endl;
-    std::cerr << error << std::endl;
-    abort();
-  }
+  if (!target)
+    throw new std::runtime_error("Could not create LLVM target " + error);
 
   TargetOptions opt;
   auto RM = Reloc::Model();
@@ -714,12 +688,10 @@ int CodegenLLVM::compile(bool debug, std::ostream &out)
     module_->print(llvm_ostream, nullptr, false, true);
   }
 
-  EngineBuilder builder(move(module_));
-  builder.setMCJITMemoryManager(std::make_unique<BPFtraceMemoryManager>(bpftrace_.sections_));
-  ee_ = std::unique_ptr<ExecutionEngine>(builder.create());
-  ee_->finalizeObject();
+  auto bpforc = std::make_unique<BpfOrc>(targetMachine);
+  bpforc->compileModule(move(module_));
 
-  return 0;
+  return move(bpforc);
 }
 
 } // namespace ast

src/ast/codegen_llvm.h

@@ -52,7 +52,7 @@ public:
 
   void createLog2Function();
   void createStrcmpFunction();
-  int compile(bool debug=false, std::ostream &out=std::cerr);
+  std::unique_ptr<BpfOrc> compile(bool debug=false, std::ostream &out=std::cerr);
 
 private:
   Node *root_;

src/attached_probe.cpp

@@ -42,7 +42,7 @@ bpf_prog_type progtype(ProbeType t)
 }
 
 
-AttachedProbe::AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> &func)
+AttachedProbe::AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func)
   : probe_(probe), func_(func)
 {
   load_prog();

src/attached_probe.h

@@ -12,7 +12,7 @@ bpf_prog_type progtype(ProbeType t);
 class AttachedProbe
 {
 public:
-  AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> &func);
+  AttachedProbe(Probe &probe, std::tuple<uint8_t *, uintptr_t> func);
   ~AttachedProbe();
   AttachedProbe(const AttachedProbe &) = delete;
   AttachedProbe& operator=(const AttachedProbe &) = delete;
@@ -29,7 +29,7 @@ private:
   void attach_profile();
 
   Probe &probe_;
-  std::tuple<uint8_t *, uintptr_t> &func_;
+  std::tuple<uint8_t *, uintptr_t> func_;
   std::vector<int> perf_event_fds_;
   int progfd_;
 };

src/bpforc.h

+#pragma once
+
+#include "llvm/ExecutionEngine/ExecutionEngine.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/SectionMemoryManager.h"
+#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
+#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
+#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
+#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/Target/TargetMachine.h"
+
+namespace bpftrace {
+
+using namespace llvm;
+using namespace llvm::orc;
+
+class MemoryManager : public SectionMemoryManager
+{
+public:
+  explicit MemoryManager(std::map<std::string, std::tuple<uint8_t *, uintptr_t>> &sections)
+    : sections_(sections) { }
+  uint8_t *allocateCodeSection(uintptr_t Size, unsigned Alignment, unsigned SectionID, StringRef SectionName) override
+  {
+    uint8_t *addr = SectionMemoryManager::allocateCodeSection(Size, Alignment, SectionID, SectionName);
+    sections_[SectionName.str()] = std::make_tuple(addr, Size);
+    return addr;
+  }
+
+  uint8_t *allocateDataSection(uintptr_t Size, unsigned Alignment, unsigned SectionID, StringRef SectionName, bool isReadOnly) override
+  {
+    uint8_t *addr = SectionMemoryManager::allocateDataSection(Size, Alignment, SectionID, SectionName, isReadOnly);
+    sections_[SectionName.str()] = std::make_tuple(addr, Size);
+    return addr;
+  }
+
+private:
+  std::map<std::string, std::tuple<uint8_t *, uintptr_t>> &sections_;
+};
+
+class BpfOrc
+{
+private:
+  std::unique_ptr<TargetMachine> TM;
+  RTDyldObjectLinkingLayer ObjectLayer;
+  IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
+
+public:
+  std::map<std::string, std::tuple<uint8_t *, uintptr_t>> sections_;
+
+  using ModuleHandle = decltype(CompileLayer)::ModuleHandleT;
+
+  BpfOrc(TargetMachine *TM_)
+    : TM(TM_),
+      ObjectLayer([this]() { return std::make_shared<MemoryManager>(sections_); }),
+      CompileLayer(ObjectLayer, SimpleCompiler(*TM))
+  {
+  }
+
+  void compileModule(std::unique_ptr<Module> M)
+  {
+    auto mod = addModule(move(M));
+    CompileLayer.emitAndFinalize(mod);
+  }
+
+  ModuleHandle addModule(std::unique_ptr<Module> M) {
+    // We don't actually care about resolving symbols from other modules
+    auto Resolver = createLambdaResolver(
+        [](const std::string &Name) { return JITSymbol(nullptr); },
+        [](const std::string &Name) { return JITSymbol(nullptr); });
+
+    return cantFail(CompileLayer.addModule(std::move(M), std::move(Resolver)));
+  }
+};
+
+} // namespace bpftrace

src/bpftrace.cpp

@@ -9,6 +9,7 @@
 #include "bcc_syms.h"
 #include "perf_reader.h"
 
+#include "bpforc.h"
 #include "bpftrace.h"
 #include "attached_probe.h"
 #include "triggers.h"
@@ -195,10 +196,10 @@ void perf_event_lost(void *cb_cookie, uint64_t lost)
   printf("Lost %lu events\n", lost);
 }
 
-std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe)
+std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe, const BpfOrc &bpforc)
 {
-  auto func = sections_.find("s_" + probe.prog_name);
-  if (func == sections_.end())
+  auto func = bpforc.sections_.find("s_" + probe.prog_name);
+  if (func == bpforc.sections_.end())
   {
     std::cerr << "Code not generated for probe: " << probe.name << std::endl;
     return nullptr;
@@ -214,11 +215,11 @@ std::unique_ptr<AttachedProbe> BPFtrace::attach_probe(Probe &probe)
   return nullptr;
 }
 
-int BPFtrace::run()
+int BPFtrace::run(std::unique_ptr<BpfOrc> bpforc)
 {
   for (Probe &probe : special_probes_)
   {
-    auto attached_probe = attach_probe(probe);
+    auto attached_probe = attach_probe(probe, *bpforc.get());
     if (attached_probe == nullptr)
       return -1;
     special_attached_probes_.push_back(std::move(attached_probe));
@@ -232,7 +233,7 @@ int BPFtrace::run()
 
   for (Probe &probe : probes_)
   {
-    auto attached_probe = attach_probe(probe);
+    auto attached_probe = attach_probe(probe, *bpforc.get());
     if (attached_probe == nullptr)
       return -1;
     attached_probes_.push_back(std::move(attached_probe));

src/bpftrace.h

@@ -16,6 +16,8 @@
 
 namespace bpftrace {
 
+class BpfOrc;
+
 class BPFtrace
 {
 public:
@@ -23,14 +25,13 @@ public:
   virtual ~BPFtrace() { }
   virtual int add_probe(ast::Probe &p);
   int num_probes() const;
-  int run();
+  int run(std::unique_ptr<BpfOrc> bpforc);
   int print_maps();
   std::string get_stack(uint32_t stackid, bool ustack, int indent=0);
   std::string resolve_sym(uintptr_t addr, bool show_offset=false);
   std::string resolve_usym(uintptr_t addr) const;
 
   std::map<std::string, std::unique_ptr<IMap>> maps_;
-  std::map<std::string, std::tuple<uint8_t *, uintptr_t>> sections_;
   std::map<std::string, Struct> structs_;
   std::vector<std::tuple<std::string, std::vector<SizedType>>> printf_args_;
   std::unique_ptr<IMap> stackid_map_;
@@ -51,7 +52,7 @@ private:
   int ncpus_;
   int online_cpus_;
 
-  std::unique_ptr<AttachedProbe> attach_probe(Probe &probe);
+  std::unique_ptr<AttachedProbe> attach_probe(Probe &probe, const BpfOrc &bpforc);
   int setup_perf_events();
   void poll_perf_events(int epollfd, int timeout=-1);
   int print_map(IMap &map);

src/main.cpp

 #include <iostream>
 #include <signal.h>
 
+#include "bpforc.h"
 #include "bpftrace.h"
 #include "codegen_llvm.h"
 #include "driver.h"
@@ -83,9 +84,7 @@ int main(int argc, char *argv[])
     return err;
 
   ast::CodegenLLVM llvm(driver.root_, bpftrace);
-  err = llvm.compile(debug);
-  if (err)
-    return err;
+  auto bpforc = llvm.compile(debug);
 
   if (debug)
     return 0;
@@ -106,7 +105,7 @@ int main(int argc, char *argv[])
   else
     std::cout << "Attaching " << bpftrace.num_probes() << " probes..." << std::endl;
 
-  err = bpftrace.run();
+  err = bpftrace.run(move(bpforc));
   if (err)
     return err;
 

tests/codegen.cpp

 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "bpforc.h"
 #include "bpftrace.h"
 #include "codegen_llvm.h"
 #include "driver.h"
@@ -12,6 +13,24 @@ namespace codegen {
 
 using ::testing::_;
 
+TEST(codegen, populate_sections)
+{
+  BPFtrace bpftrace;
+  Driver driver;
+
+  ASSERT_EQ(driver.parse_str("kprobe:foo { 1 } kprobe:bar { 1 }"), 0);
+  ast::SemanticAnalyser semantics(driver.root_, bpftrace);
+  ASSERT_EQ(semantics.analyse(), 0);
+  std::stringstream out;
+  ast::CodegenLLVM codegen(driver.root_, bpftrace);
+  auto bpforc = codegen.compile(true, out);
+
+  // Check sections are populated
+  ASSERT_EQ(bpforc->sections_.size(), 2);
+  ASSERT_EQ(bpforc->sections_.count("s_kprobe:foo"), 1);
+  ASSERT_EQ(bpforc->sections_.count("s_kprobe:bar"), 1);
+}
+
 std::string header = R"HEAD(; ModuleID = 'bpftrace'
 source_filename = "bpftrace"
 target datalayout = "e-m:e-p:64:64-i64:64-n32:64-S128"
@@ -33,7 +52,7 @@ void test(const std::string &input, const std::string expected_output)
 
   std::stringstream out;
   ast::CodegenLLVM codegen(driver.root_, bpftrace);
-  ASSERT_EQ(codegen.compile(true, out), 0);
+  codegen.compile(true, out);
 
   std::string full_expected_output = header + expected_output;
   EXPECT_EQ(full_expected_output, out.str());