1//===- JIT.cpp - Target independent JIT infrastructure --------------------===//
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//===----------------------------------------------------------------------===//
10
11#include "JIT.h"
12
13#include "Shared/Debug.h"
14#include "Shared/Utils.h"
15
16#include "PluginInterface.h"
17#include "omptarget.h"
18
19#include "llvm/ADT/SmallVector.h"
20#include "llvm/CodeGen/CommandFlags.h"
21#include "llvm/CodeGen/MachineModuleInfo.h"
22#include "llvm/IR/LLVMContext.h"
23#include "llvm/IR/LLVMRemarkStreamer.h"
24#include "llvm/IR/LegacyPassManager.h"
25#include "llvm/IRReader/IRReader.h"
26#include "llvm/InitializePasses.h"
27#include "llvm/MC/TargetRegistry.h"
28#include "llvm/Object/IRObjectFile.h"
29#include "llvm/Passes/OptimizationLevel.h"
30#include "llvm/Passes/PassBuilder.h"
31#include "llvm/Support/MemoryBuffer.h"
32#include "llvm/Support/SourceMgr.h"
33#include "llvm/Support/TargetSelect.h"
34#include "llvm/Support/TimeProfiler.h"
35#include "llvm/Support/ToolOutputFile.h"
36#include "llvm/Support/raw_ostream.h"
37#include "llvm/Target/TargetMachine.h"
38#include "llvm/Target/TargetOptions.h"
39#include "llvm/TargetParser/SubtargetFeature.h"
40
41#include <mutex>
42#include <shared_mutex>
43#include <system_error>
44
45using namespace llvm;
46using namespace llvm::object;
47using namespace omp;
48using namespace omp::target;
49
50namespace {
51
52bool isImageBitcode(const __tgt_device_image &Image) {
53 StringRef Binary(reinterpret_cast<const char *>(Image.ImageStart),
54 target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
55
56 return identify_magic(magic: Binary) == file_magic::bitcode;
57}
58
59std::once_flag InitFlag;
60
61void init(Triple TT) {
62 codegen::RegisterCodeGenFlags();
63#ifdef LIBOMPTARGET_JIT_NVPTX
64 if (TT.isNVPTX()) {
65 LLVMInitializeNVPTXTargetInfo();
66 LLVMInitializeNVPTXTarget();
67 LLVMInitializeNVPTXTargetMC();
68 LLVMInitializeNVPTXAsmPrinter();
69 }
70#endif
71#ifdef LIBOMPTARGET_JIT_AMDGPU
72 if (TT.isAMDGPU()) {
73 LLVMInitializeAMDGPUTargetInfo();
74 LLVMInitializeAMDGPUTarget();
75 LLVMInitializeAMDGPUTargetMC();
76 LLVMInitializeAMDGPUAsmPrinter();
77 }
78#endif
79}
80
81Expected<std::unique_ptr<Module>>
82createModuleFromMemoryBuffer(std::unique_ptr<MemoryBuffer> &MB,
83 LLVMContext &Context) {
84 SMDiagnostic Err;
85 auto Mod = parseIR(*MB, Err, Context);
86 if (!Mod)
87 return make_error<StringError>(Args: "Failed to create module",
88 Args: inconvertibleErrorCode());
89 return std::move(Mod);
90}
91Expected<std::unique_ptr<Module>>
92createModuleFromImage(const __tgt_device_image &Image, LLVMContext &Context) {
93 StringRef Data((const char *)Image.ImageStart,
94 target::getPtrDiff(Image.ImageEnd, Image.ImageStart));
95 std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer(
96 InputData: Data, /*BufferName=*/"", /*RequiresNullTerminator=*/false);
97 return createModuleFromMemoryBuffer(MB, Context);
98}
99
100OptimizationLevel getOptLevel(unsigned OptLevel) {
101 switch (OptLevel) {
102 case 0:
103 return OptimizationLevel::O0;
104 case 1:
105 return OptimizationLevel::O1;
106 case 2:
107 return OptimizationLevel::O2;
108 case 3:
109 return OptimizationLevel::O3;
110 }
111 llvm_unreachable("Invalid optimization level");
112}
113
114Expected<std::unique_ptr<TargetMachine>>
115createTargetMachine(Module &M, std::string CPU, unsigned OptLevel) {
116 Triple TT(M.getTargetTriple());
117 std::optional<CodeGenOptLevel> CGOptLevelOrNone =
118 CodeGenOpt::getLevel(OptLevel);
119 assert(CGOptLevelOrNone && "Invalid optimization level");
120 CodeGenOptLevel CGOptLevel = *CGOptLevelOrNone;
121
122 std::string Msg;
123 const Target *T = TargetRegistry::lookupTarget(Triple: M.getTargetTriple(), Error&: Msg);
124 if (!T)
125 return make_error<StringError>(Args&: Msg, Args: inconvertibleErrorCode());
126
127 SubtargetFeatures Features;
128 Features.getDefaultSubtargetFeatures(Triple: TT);
129
130 std::optional<Reloc::Model> RelocModel;
131 if (M.getModuleFlag("PIC Level"))
132 RelocModel =
133 M.getPICLevel() == PICLevel::NotPIC ? Reloc::Static : Reloc::PIC_;
134
135 std::optional<CodeModel::Model> CodeModel = M.getCodeModel();
136
137 TargetOptions Options = codegen::InitTargetOptionsFromCodeGenFlags(TheTriple: TT);
138
139 std::unique_ptr<TargetMachine> TM(
140 T->createTargetMachine(M.getTargetTriple(), CPU, Features.getString(),
141 Options, RelocModel, CodeModel, CGOptLevel));
142 if (!TM)
143 return make_error<StringError>(Args: "Failed to create target machine",
144 Args: inconvertibleErrorCode());
145 return std::move(TM);
146}
147
148} // namespace
149
150JITEngine::JITEngine(Triple::ArchType TA) : TT(Triple::getArchTypeName(TA)) {
151 std::call_once(InitFlag, init, TT);
152}
153
154void JITEngine::opt(TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
155 unsigned OptLevel) {
156 PipelineTuningOptions PTO;
157 std::optional<PGOOptions> PGOOpt;
158
159 LoopAnalysisManager LAM;
160 FunctionAnalysisManager FAM;
161 CGSCCAnalysisManager CGAM;
162 ModuleAnalysisManager MAM;
163 ModulePassManager MPM;
164
165 PassBuilder PB(TM, PTO, PGOOpt, nullptr);
166
167 FAM.registerPass([&] { return TargetLibraryAnalysis(*TLII); });
168
169 // Register all the basic analyses with the managers.
170 PB.registerModuleAnalyses(MAM);
171 PB.registerCGSCCAnalyses(CGAM);
172 PB.registerFunctionAnalyses(FAM);
173 PB.registerLoopAnalyses(LAM);
174 PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
175
176 MPM.addPass(PB.buildPerModuleDefaultPipeline(getOptLevel(OptLevel)));
177 MPM.run(M, MAM);
178}
179
180void JITEngine::codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII,
181 Module &M, raw_pwrite_stream &OS) {
182 legacy::PassManager PM;
183 PM.add(new TargetLibraryInfoWrapperPass(*TLII));
184 MachineModuleInfoWrapperPass *MMIWP = new MachineModuleInfoWrapperPass(
185 reinterpret_cast<LLVMTargetMachine *>(TM));
186 TM->addPassesToEmitFile(PM, OS, nullptr,
187 TT.isNVPTX() ? CodeGenFileType::AssemblyFile
188 : CodeGenFileType::ObjectFile,
189 /*DisableVerify=*/false, MMIWP);
190
191 PM.run(M);
192}
193
194Expected<std::unique_ptr<MemoryBuffer>>
195JITEngine::backend(Module &M, const std::string &ComputeUnitKind,
196 unsigned OptLevel) {
197
198 auto RemarksFileOrErr = setupLLVMOptimizationRemarks(
199 M.getContext(), /*RemarksFilename=*/"", /*RemarksPasses=*/"",
200 /*RemarksFormat=*/"", /*RemarksWithHotness=*/false);
201 if (Error E = RemarksFileOrErr.takeError())
202 return std::move(E);
203 if (*RemarksFileOrErr)
204 (*RemarksFileOrErr)->keep();
205
206 auto TMOrErr = createTargetMachine(M, ComputeUnitKind, OptLevel);
207 if (!TMOrErr)
208 return TMOrErr.takeError();
209
210 std::unique_ptr<TargetMachine> TM = std::move(*TMOrErr);
211 TargetLibraryInfoImpl TLII(TT);
212
213 if (PreOptIRModuleFileName.isPresent()) {
214 std::error_code EC;
215 raw_fd_stream FD(PreOptIRModuleFileName.get(), EC);
216 if (EC)
217 return createStringError(
218 EC, "Could not open %s to write the pre-opt IR module\n",
219 PreOptIRModuleFileName.get().c_str());
220 M.print(FD, nullptr);
221 }
222
223 if (!JITSkipOpt)
224 opt(TM.get(), &TLII, M, OptLevel);
225
226 if (PostOptIRModuleFileName.isPresent()) {
227 std::error_code EC;
228 raw_fd_stream FD(PostOptIRModuleFileName.get(), EC);
229 if (EC)
230 return createStringError(
231 EC, "Could not open %s to write the post-opt IR module\n",
232 PreOptIRModuleFileName.get().c_str());
233 M.print(FD, nullptr);
234 }
235
236 // Prepare the output buffer and stream for codegen.
237 SmallVector<char> CGOutputBuffer;
238 raw_svector_ostream OS(CGOutputBuffer);
239
240 codegen(TM.get(), &TLII, M, OS);
241
242 return MemoryBuffer::getMemBufferCopy(OS.str());
243}
244
245Expected<std::unique_ptr<MemoryBuffer>>
246JITEngine::getOrCreateObjFile(const __tgt_device_image &Image, LLVMContext &Ctx,
247 const std::string &ComputeUnitKind) {
248
249 // Check if the user replaces the module at runtime with a finished object.
250 if (ReplacementObjectFileName.isPresent()) {
251 auto MBOrErr =
252 MemoryBuffer::getFileOrSTDIN(ReplacementObjectFileName.get());
253 if (!MBOrErr)
254 return createStringError(MBOrErr.getError(),
255 "Could not read replacement obj from %s\n",
256 ReplacementModuleFileName.get().c_str());
257 return std::move(*MBOrErr);
258 }
259
260 Module *Mod = nullptr;
261 // Check if the user replaces the module at runtime or we read it from the
262 // image.
263 // TODO: Allow the user to specify images per device (Arch + ComputeUnitKind).
264 if (!ReplacementModuleFileName.isPresent()) {
265 auto ModOrErr = createModuleFromImage(Image, Ctx);
266 if (!ModOrErr)
267 return ModOrErr.takeError();
268 Mod = ModOrErr->release();
269 } else {
270 auto MBOrErr =
271 MemoryBuffer::getFileOrSTDIN(ReplacementModuleFileName.get());
272 if (!MBOrErr)
273 return createStringError(MBOrErr.getError(),
274 "Could not read replacement module from %s\n",
275 ReplacementModuleFileName.get().c_str());
276 auto ModOrErr = createModuleFromMemoryBuffer(MBOrErr.get(), Ctx);
277 if (!ModOrErr)
278 return ModOrErr.takeError();
279 Mod = ModOrErr->release();
280 }
281
282 return backend(*Mod, ComputeUnitKind, JITOptLevel);
283}
284
285Expected<const __tgt_device_image *>
286JITEngine::compile(const __tgt_device_image &Image,
287 const std::string &ComputeUnitKind,
288 PostProcessingFn PostProcessing) {
289 std::lock_guard<std::mutex> Lock(ComputeUnitMapMutex);
290
291 // Check if we JITed this image for the given compute unit kind before.
292 ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];
293 if (__tgt_device_image *JITedImage = CUI.TgtImageMap.lookup(&Image))
294 return JITedImage;
295
296 auto ObjMBOrErr = getOrCreateObjFile(Image, CUI.Context, ComputeUnitKind);
297 if (!ObjMBOrErr)
298 return ObjMBOrErr.takeError();
299
300 auto ImageMBOrErr = PostProcessing(std::move(*ObjMBOrErr));
301 if (!ImageMBOrErr)
302 return ImageMBOrErr.takeError();
303
304 CUI.JITImages.push_back(std::move(*ImageMBOrErr));
305 __tgt_device_image *&JITedImage = CUI.TgtImageMap[&Image];
306 JITedImage = new __tgt_device_image();
307 *JITedImage = Image;
308
309 auto &ImageMB = CUI.JITImages.back();
310
311 JITedImage->ImageStart = const_cast<char *>(ImageMB->getBufferStart());
312 JITedImage->ImageEnd = const_cast<char *>(ImageMB->getBufferEnd());
313
314 return JITedImage;
315}
316
317Expected<const __tgt_device_image *>
318JITEngine::process(const __tgt_device_image &Image,
319 target::plugin::GenericDeviceTy &Device) {
320 const std::string &ComputeUnitKind = Device.getComputeUnitKind();
321
322 PostProcessingFn PostProcessing = [&Device](std::unique_ptr<MemoryBuffer> MB)
323 -> Expected<std::unique_ptr<MemoryBuffer>> {
324 return Device.doJITPostProcessing(std::move(MB));
325 };
326
327 if (isImageBitcode(Image))
328 return compile(Image, ComputeUnitKind, PostProcessing);
329
330 return &Image;
331}
332
333Expected<bool> JITEngine::checkBitcodeImage(StringRef Buffer) const {
334 TimeTraceScope TimeScope("Check bitcode image");
335
336 assert(identify_magic(Buffer) == file_magic::bitcode &&
337 "Input is not bitcode");
338
339 LLVMContext Context;
340 auto ModuleOrErr = getLazyBitcodeModule(MemoryBufferRef(Buffer, ""), Context,
341 /*ShouldLazyLoadMetadata=*/true);
342 if (!ModuleOrErr)
343 return ModuleOrErr.takeError();
344 Module &M = **ModuleOrErr;
345
346 return Triple(M.getTargetTriple()).getArch() == TT.getArch();
347}
348

source code of offload/plugins-nextgen/common/src/JIT.cpp