1 //===--- LLJITWithThinLTOSummaries.cpp - Module summaries as LLJIT input --===//
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 // In this example we will use a module summary index file produced for ThinLTO
10 // to (A) find the module that defines the main entry point and (B) find all
11 // extra modules that we need. We will do this in five steps:
12 //
13 // (1) Read the index file and parse the module summary index.
14 // (2) Find the path of the module that defines "main".
15 // (3) Parse the main module and create a matching LLJIT.
16 // (4) Add all modules to the LLJIT that are covered by the index.
17 // (5) Look up and run the JIT'd function.
18 //
19 // The index file name must be passed in as command line argument. Please find
20 // this test for instructions on creating the index file:
21 //
22 //       llvm/test/Examples/OrcV2Examples/lljit-with-thinlto-summaries.test
23 //
24 // If you use "build" as the build directory, you can run the test from the root
25 // of the monorepo like this:
26 //
27 // > build/bin/llvm-lit -a \
28 //       llvm/test/Examples/OrcV2Examples/lljit-with-thinlto-summaries.test
29 //
30 //===----------------------------------------------------------------------===//
31 
32 #include "llvm/ADT/STLExtras.h"
33 #include "llvm/ADT/StringRef.h"
34 #include "llvm/Bitcode/BitcodeReader.h"
35 #include "llvm/ExecutionEngine/Orc/Core.h"
36 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
37 #include "llvm/ExecutionEngine/Orc/LLJIT.h"
38 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
39 #include "llvm/IR/GlobalValue.h"
40 #include "llvm/IR/LLVMContext.h"
41 #include "llvm/IR/ModuleSummaryIndex.h"
42 #include "llvm/Support/CommandLine.h"
43 #include "llvm/Support/Error.h"
44 #include "llvm/Support/InitLLVM.h"
45 #include "llvm/Support/MemoryBuffer.h"
46 #include "llvm/Support/TargetSelect.h"
47 #include "llvm/Support/raw_ostream.h"
48 
49 #include <string>
50 #include <system_error>
51 #include <vector>
52 
53 using namespace llvm;
54 using namespace llvm::orc;
55 
56 // Path of the module summary index file.
57 cl::opt<std::string> IndexFile{cl::desc("<module summary index>"),
58                                cl::Positional, cl::init("-")};
59 
60 // Describe a fail state that is caused by the given ModuleSummaryIndex
61 // providing multiple definitions of the given global value name. It will dump
62 // name and GUID for the global value and list the paths of the modules covered
63 // by the index.
64 class DuplicateDefinitionInSummary
65     : public ErrorInfo<DuplicateDefinitionInSummary> {
66 public:
67   static char ID;
68 
69   DuplicateDefinitionInSummary(std::string GlobalValueName, ValueInfo VI)
70       : GlobalValueName(std::move(GlobalValueName)) {
71     ModulePaths.reserve(VI.getSummaryList().size());
72     for (const auto &S : VI.getSummaryList())
73       ModulePaths.push_back(S->modulePath().str());
74     llvm::sort(ModulePaths);
75   }
76 
77   void log(raw_ostream &OS) const override {
78     OS << "Duplicate symbol for global value '" << GlobalValueName
79        << "' (GUID: " << GlobalValue::getGUID(GlobalValueName) << ") in:\n";
80     for (const std::string &Path : ModulePaths) {
81       OS << "    " << Path << "\n";
82     }
83   }
84 
85   std::error_code convertToErrorCode() const override {
86     return inconvertibleErrorCode();
87   }
88 
89 private:
90   std::string GlobalValueName;
91   std::vector<std::string> ModulePaths;
92 };
93 
94 // Describe a fail state where the given global value name was not found in the
95 // given ModuleSummaryIndex. It will dump name and GUID for the global value and
96 // list the paths of the modules covered by the index.
97 class DefinitionNotFoundInSummary
98     : public ErrorInfo<DefinitionNotFoundInSummary> {
99 public:
100   static char ID;
101 
102   DefinitionNotFoundInSummary(std::string GlobalValueName,
103                               ModuleSummaryIndex &Index)
104       : GlobalValueName(std::move(GlobalValueName)) {
105     ModulePaths.reserve(Index.modulePaths().size());
106     for (const auto &Entry : Index.modulePaths())
107       ModulePaths.push_back(Entry.first().str());
108     llvm::sort(ModulePaths);
109   }
110 
111   void log(raw_ostream &OS) const override {
112     OS << "No symbol for global value '" << GlobalValueName
113        << "' (GUID: " << GlobalValue::getGUID(GlobalValueName) << ") in:\n";
114     for (const std::string &Path : ModulePaths) {
115       OS << "    " << Path << "\n";
116     }
117   }
118 
119   std::error_code convertToErrorCode() const override {
120     return llvm::inconvertibleErrorCode();
121   }
122 
123 private:
124   std::string GlobalValueName;
125   std::vector<std::string> ModulePaths;
126 };
127 
128 char DuplicateDefinitionInSummary::ID = 0;
129 char DefinitionNotFoundInSummary::ID = 0;
130 
131 // Lookup the a function in the ModuleSummaryIndex and return the path of the
132 // module that defines it. Paths in the ModuleSummaryIndex are relative to the
133 // build directory of the covered modules.
134 Expected<StringRef> getMainModulePath(StringRef FunctionName,
135                                       ModuleSummaryIndex &Index) {
136   // Summaries use unmangled names.
137   GlobalValue::GUID G = GlobalValue::getGUID(FunctionName);
138   ValueInfo VI = Index.getValueInfo(G);
139 
140   // We need a unique definition, otherwise don't try further.
141   if (!VI || VI.getSummaryList().empty())
142     return make_error<DefinitionNotFoundInSummary>(FunctionName.str(), Index);
143   if (VI.getSummaryList().size() > 1)
144     return make_error<DuplicateDefinitionInSummary>(FunctionName.str(), VI);
145 
146   GlobalValueSummary *S = VI.getSummaryList().front()->getBaseObject();
147   if (!isa<FunctionSummary>(S))
148     return createStringError(inconvertibleErrorCode(),
149                              "Entry point is not a function: " + FunctionName);
150 
151   // Return a reference. ModuleSummaryIndex owns the module paths.
152   return S->modulePath();
153 }
154 
155 // Parse the bitcode module from the given path into a ThreadSafeModule.
156 Expected<ThreadSafeModule> loadModule(StringRef Path,
157                                       orc::ThreadSafeContext TSCtx) {
158   outs() << "About to load module: " << Path << "\n";
159 
160   Expected<std::unique_ptr<MemoryBuffer>> BitcodeBuffer =
161       errorOrToExpected(MemoryBuffer::getFile(Path));
162   if (!BitcodeBuffer)
163     return BitcodeBuffer.takeError();
164 
165   MemoryBufferRef BitcodeBufferRef = (**BitcodeBuffer).getMemBufferRef();
166   Expected<std::unique_ptr<Module>> M =
167       parseBitcodeFile(BitcodeBufferRef, *TSCtx.getContext());
168   if (!M)
169     return M.takeError();
170 
171   return ThreadSafeModule(std::move(*M), std::move(TSCtx));
172 }
173 
174 int main(int Argc, char *Argv[]) {
175   InitLLVM X(Argc, Argv);
176 
177   InitializeNativeTarget();
178   InitializeNativeTargetAsmPrinter();
179 
180   cl::ParseCommandLineOptions(Argc, Argv, "LLJITWithThinLTOSummaries");
181 
182   ExitOnError ExitOnErr;
183   ExitOnErr.setBanner(std::string(Argv[0]) + ": ");
184 
185   // (1) Read the index file and parse the module summary index.
186   std::unique_ptr<MemoryBuffer> SummaryBuffer =
187       ExitOnErr(errorOrToExpected(MemoryBuffer::getFile(IndexFile)));
188 
189   std::unique_ptr<ModuleSummaryIndex> SummaryIndex =
190       ExitOnErr(getModuleSummaryIndex(SummaryBuffer->getMemBufferRef()));
191 
192   // (2) Find the path of the module that defines "main".
193   std::string MainFunctionName = "main";
194   StringRef MainModulePath =
195       ExitOnErr(getMainModulePath(MainFunctionName, *SummaryIndex));
196 
197   // (3) Parse the main module and create a matching LLJIT.
198   ThreadSafeContext TSCtx(std::make_unique<LLVMContext>());
199   ThreadSafeModule MainModule = ExitOnErr(loadModule(MainModulePath, TSCtx));
200 
201   auto Builder = LLJITBuilder();
202 
203   MainModule.withModuleDo([&](Module &M) {
204     if (M.getTargetTriple().empty()) {
205       Builder.setJITTargetMachineBuilder(
206           ExitOnErr(JITTargetMachineBuilder::detectHost()));
207     } else {
208       Builder.setJITTargetMachineBuilder(
209           JITTargetMachineBuilder(Triple(M.getTargetTriple())));
210     }
211     if (!M.getDataLayout().getStringRepresentation().empty())
212       Builder.setDataLayout(M.getDataLayout());
213   });
214 
215   auto J = ExitOnErr(Builder.create());
216 
217   // (4) Add all modules to the LLJIT that are covered by the index.
218   JITDylib &JD = J->getMainJITDylib();
219 
220   for (const auto &Entry : SummaryIndex->modulePaths()) {
221     StringRef Path = Entry.first();
222     ThreadSafeModule M = (Path == MainModulePath)
223                              ? std::move(MainModule)
224                              : ExitOnErr(loadModule(Path, TSCtx));
225     ExitOnErr(J->addIRModule(JD, std::move(M)));
226   }
227 
228   // (5) Look up and run the JIT'd function.
229   auto MainSym = ExitOnErr(J->lookup(MainFunctionName));
230 
231   using MainFnPtr = int (*)(int, char *[]);
232   MainFnPtr MainFunction =
233       jitTargetAddressToFunction<MainFnPtr>(MainSym.getAddress());
234 
235   int Result = runAsMain(MainFunction, {}, MainModulePath);
236   outs() << "'" << MainFunctionName << "' finished with exit code: " << Result
237          << "\n";
238 
239   return 0;
240 }
241