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