llvm/examples/OrcV2Examples/LLJITWithObjectLinkingLayerPlugin/LLJITWithObjectLinkingLayerPlugin.cpp

   1 //===--------------- LLJITWithCustomObjectLinkingLayer.cpp ----------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file shows how to switch LLJIT to use a custom object linking layer (we
  10 // use ObjectLinkingLayer, which is backed by JITLink, as an example).
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #include "llvm/ADT/StringMap.h"
  15 #include "llvm/ExecutionEngine/JITLink/JITLink.h"
  16 #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
  17 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
  18 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
  19 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
  20 #include "llvm/Support/InitLLVM.h"
  21 #include "llvm/Support/TargetSelect.h"
  22 #include "llvm/Support/raw_ostream.h"
  23
  24 #include "../ExampleModules.h"
  25
  26 using namespace llvm;
  27 using namespace llvm::orc;
  28
  29 ExitOnError ExitOnErr;
  30
  31 const llvm::StringRef TestMod =
  32     R"(
  33   define i32 @callee() {
  34   entry:
  35     ret i32 7
  36   }
  37
  38   define i32 @entry() {
  39   entry:
  40     %0 = call i32 @callee()
  41     ret i32 %0
  42   }
  43 )";
  44
  45 class MyPlugin : public ObjectLinkingLayer::Plugin {
  46 public:
  47   // The modifyPassConfig callback gives us a chance to inspect the
  48   // MaterializationResponsibility and target triple for the object being
  49   // linked, then add any JITLink passes that we would like to run on the
  50   // link graph. A pass is just a function object that is callable as
  51   // Error(jitlink::LinkGraph&). In this case we will add two passes
  52   // defined as lambdas that call the printLinkerGraph method on our
  53   // plugin: One to run before the linker applies fixups and another to
  54   // run afterwards.
  55   void modifyPassConfig(MaterializationResponsibility &MR,
  56                         jitlink::LinkGraph &LG,
  57                         jitlink::PassConfiguration &Config) override {
  58
  59     outs() << "MyPlugin -- Modifying pass config for " << LG.getName() << " ("
  60            << LG.getTargetTriple().str() << "):\n";
  61
  62     // Print sections, symbol names and addresses, and any edges for the
  63     // associated blocks at the 'PostPrune' phase of JITLink (after
  64     // dead-stripping, but before addresses are allocated in the target
  65     // address space. See llvm/docs/JITLink.rst).
  66     //
  67     // Experiment with adding the 'printGraph' pass at other points in the
  68     // pipeline. E.g. PrePrunePasses, PostAllocationPasses, and
  69     // PostFixupPasses.
  70     Config.PostPrunePasses.push_back(printGraph);
  71   }
  72
  73   void notifyLoaded(MaterializationResponsibility &MR) override {
  74     outs() << "Loading object defining " << MR.getSymbols() << "\n";
  75   }
  76
  77   Error notifyEmitted(MaterializationResponsibility &MR) override {
  78     outs() << "Emitted object defining " << MR.getSymbols() << "\n";
  79     return Error::success();
  80   }
  81
  82   Error notifyFailed(MaterializationResponsibility &MR) override {
  83     return Error::success();
  84   }
  85
  86   Error notifyRemovingResources(ResourceKey K) override {
  87     return Error::success();
  88   }
  89
  90   void notifyTransferringResources(ResourceKey DstKey,
  91                                    ResourceKey SrcKey) override {}
  92
  93 private:
  94   static void printBlockContent(jitlink::Block &B) {
  95     constexpr JITTargetAddress LineWidth = 16;
  96
  97     if (B.isZeroFill()) {
  98       outs() << "    " << formatv("{0:x16}", B.getAddress()) << ": "
  99              << B.getSize() << " bytes of zero-fill.\n";
 100       return;
 101     }
 102
 103     JITTargetAddress InitAddr = B.getAddress() & ~(LineWidth - 1);
 104     JITTargetAddress StartAddr = B.getAddress();
 105     JITTargetAddress EndAddr = B.getAddress() + B.getSize();
 106     auto *Data = reinterpret_cast<const uint8_t *>(B.getContent().data());
 107
 108     for (JITTargetAddress CurAddr = InitAddr; CurAddr != EndAddr; ++CurAddr) {
 109       if (CurAddr % LineWidth == 0)
 110         outs() << "          " << formatv("{0:x16}", CurAddr) << ": ";
 111       if (CurAddr < StartAddr)
 112         outs() << "   ";
 113       else
 114         outs() << formatv("{0:x-2}", Data[CurAddr - StartAddr]) << " ";
 115       if (CurAddr % LineWidth == LineWidth - 1)
 116         outs() << "\n";
 117     }
 118     if (EndAddr % LineWidth != 0)
 119       outs() << "\n";
 120   }
 121
 122   static Error printGraph(jitlink::LinkGraph &G) {
 123
 124     DenseSet<jitlink::Block *> BlocksAlreadyVisited;
 125
 126     outs() << "Graph \"" << G.getName() << "\"\n";
 127     // Loop over all sections...
 128     for (auto &S : G.sections()) {
 129       outs() << "  Section " << S.getName() << ":\n";
 130
 131       // Loop over all symbols in the current section...
 132       for (auto *Sym : S.symbols()) {
 133
 134         // Print the symbol's address.
 135         outs() << "    " << formatv("{0:x16}", Sym->getAddress()) << ": ";
 136
 137         // Print the symbol's name, or "<anonymous symbol>" if it doesn't have
 138         // one.
 139         if (Sym->hasName())
 140           outs() << Sym->getName() << "\n";
 141         else
 142           outs() << "<anonymous symbol>\n";
 143
 144         // Get the content block for this symbol.
 145         auto &B = Sym->getBlock();
 146
 147         if (BlocksAlreadyVisited.count(&B)) {
 148           outs() << "      Block " << formatv("{0:x16}", B.getAddress())
 149                  << " already printed.\n";
 150           continue;
 151         } else
 152           outs() << "      Block " << formatv("{0:x16}", B.getAddress())
 153                  << ":\n";
 154
 155         outs() << "        Content:\n";
 156         printBlockContent(B);
 157         BlocksAlreadyVisited.insert(&B);
 158
 159         if (!llvm::empty(B.edges())) {
 160           outs() << "        Edges:\n";
 161           for (auto &E : B.edges()) {
 162             outs() << "          "
 163                    << formatv("{0:x16}", B.getAddress() + E.getOffset())
 164                    << ": kind = " << formatv("{0:d}", E.getKind())
 165                    << ", addend = " << formatv("{0:x}", E.getAddend())
 166                    << ", target = ";
 167             jitlink::Symbol &TargetSym = E.getTarget();
 168             if (TargetSym.hasName())
 169               outs() << TargetSym.getName() << "\n";
 170             else
 171               outs() << "<anonymous target>\n";
 172           }
 173         }
 174         outs() << "\n";
 175       }
 176     }
 177     return Error::success();
 178   }
 179 };
 180
 181 static cl::opt<std::string>
 182     EntryPointName("entry", cl::desc("Symbol to call as main entry point"),
 183                    cl::init("entry"));
 184
 185 static cl::list<std::string> InputObjects(cl::Positional, cl::ZeroOrMore,
 186                                           cl::desc("input objects"));
 187
 188 int main(int argc, char *argv[]) {
 189   // Initialize LLVM.
 190   InitLLVM X(argc, argv);
 191
 192   InitializeNativeTarget();
 193   InitializeNativeTargetAsmPrinter();
 194
 195   cl::ParseCommandLineOptions(argc, argv, "LLJITWithObjectLinkingLayerPlugin");
 196   ExitOnErr.setBanner(std::string(argv[0]) + ": ");
 197
 198   // Detect the host and set code model to small.
 199   auto JTMB = ExitOnErr(JITTargetMachineBuilder::detectHost());
 200   JTMB.setCodeModel(CodeModel::Small);
 201
 202   // Create an LLJIT instance with an ObjectLinkingLayer as the base layer.
 203   // We attach our plugin in to the newly created ObjectLinkingLayer before
 204   // returning it.
 205   auto J = ExitOnErr(
 206       LLJITBuilder()
 207           .setJITTargetMachineBuilder(std::move(JTMB))
 208           .setObjectLinkingLayerCreator(
 209               [&](ExecutionSession &ES, const Triple &TT) {
 210                 // Create ObjectLinkingLayer.
 211                 auto ObjLinkingLayer = std::make_unique<ObjectLinkingLayer>(
 212                     ES, std::make_unique<jitlink::InProcessMemoryManager>());
 213                 // Add an instance of our plugin.
 214                 ObjLinkingLayer->addPlugin(std::make_unique<MyPlugin>());
 215                 return ObjLinkingLayer;
 216               })
 217           .create());
 218
 219   if (!InputObjects.empty()) {
 220
 221     // If we have input objects then reflect process symbols so the input
 222     // objects can do interesting things, like call printf.
 223     J->getMainJITDylib().addGenerator(
 224         ExitOnErr(DynamicLibrarySearchGenerator::GetForCurrentProcess(
 225             J->getDataLayout().getGlobalPrefix())));
 226
 227     // Load the input objects.
 228     for (auto InputObject : InputObjects) {
 229       auto ObjBuffer =
 230           ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(InputObject)));
 231       ExitOnErr(J->addObjectFile(std::move(ObjBuffer)));
 232     }
 233   } else {
 234     auto M = ExitOnErr(parseExampleModule(TestMod, "test-module"));
 235     M.withModuleDo([](Module &MP) {
 236       outs() << "No input objects specified. Using demo module:\n"
 237              << MP << "\n";
 238     });
 239     ExitOnErr(J->addIRModule(std::move(M)));
 240   }
 241
 242   // Look up the JIT'd function, cast it to a function pointer, then call it.
 243   auto EntrySym = ExitOnErr(J->lookup(EntryPointName));
 244   auto *Entry = (int (*)())EntrySym.getAddress();
 245
 246   int Result = Entry();
 247   outs() << "---Result---\n"
 248          << EntryPointName << "() = " << Result << "\n";
 249
 250   return 0;
 251 }