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