Expand Up
@@ -11,11 +11,11 @@
#include " JIT.h"
#include " Debug.h"
#include " PluginInterface.h"
#include " Utilities.h"
#include " omptarget.h"
#include " llvm/ADT/SmallVector.h"
#include " llvm/ADT/StringRef.h"
#include " llvm/CodeGen/CommandFlags.h"
#include " llvm/CodeGen/MachineModuleInfo.h"
#include " llvm/IR/LLVMContext.h"
Expand All
@@ -28,7 +28,6 @@
#include " llvm/Object/IRObjectFile.h"
#include " llvm/Passes/OptimizationLevel.h"
#include " llvm/Passes/PassBuilder.h"
#include " llvm/Support/Error.h"
#include " llvm/Support/MemoryBuffer.h"
#include " llvm/Support/SourceMgr.h"
#include " llvm/Support/TargetSelect.h"
Expand All
@@ -39,15 +38,23 @@
#include " llvm/Target/TargetOptions.h"
#include < mutex>
#include < shared_mutex>
#include < system_error>
using namespace llvm ;
using namespace llvm ::object;
using namespace omp ;
using namespace omp ::target;
static codegen::RegisterCodeGenFlags RCGF;
namespace {
// / A map from a bitcode image start address to its corresponding triple. If the
// / image is not in the map, it is not a bitcode image.
DenseMap<void *, Triple::ArchType> BitcodeImageMap;
std::shared_mutex BitcodeImageMapMutex;
std::once_flag InitFlag;
void init (Triple TT) {
Expand All
@@ -70,10 +77,8 @@ void init(Triple TT) {
JITTargetInitialized = true ;
}
#endif
if (!JITTargetInitialized) {
FAILURE_MESSAGE (" unsupported JIT target: %s\n " , TT.str ().c_str ());
abort ();
}
if (!JITTargetInitialized)
return ;
// Initialize passes
PassRegistry &Registry = *PassRegistry::getPassRegistry ();
Expand Down
Expand Up
@@ -125,9 +130,9 @@ createModuleFromMemoryBuffer(std::unique_ptr<MemoryBuffer> &MB,
return std::move (Mod);
}
Expected<std::unique_ptr<Module>>
createModuleFromImage (__tgt_device_image * Image, LLVMContext &Context) {
StringRef Data ((const char *)Image-> ImageStart ,
( char *) Image-> ImageEnd - ( char *) Image-> ImageStart );
createModuleFromImage (const __tgt_device_image & Image, LLVMContext &Context) {
StringRef Data ((const char *)Image. ImageStart ,
target::getPtrDiff ( Image. ImageEnd , Image. ImageStart ) );
std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer (
Data, /* BufferName */ " " , /* RequiresNullTerminator */ false );
return createModuleFromMemoryBuffer (MB, Context);
Expand Down
Expand Up
@@ -192,44 +197,11 @@ createTargetMachine(Module &M, std::string CPU, unsigned OptLevel) {
return std::move (TM);
}
// /
class JITEngine {
public:
JITEngine (Triple::ArchType TA, std::string MCpu)
: TT(Triple::getArchTypeName(TA)), CPU(MCpu),
ReplacementModuleFileName (" LIBOMPTARGET_JIT_REPLACEMENT_MODULE" ),
PreOptIRModuleFileName(" LIBOMPTARGET_JIT_PRE_OPT_IR_MODULE" ),
PostOptIRModuleFileName(" LIBOMPTARGET_JIT_POST_OPT_IR_MODULE" ) {
std::call_once (InitFlag, init, TT);
}
// / Run jit compilation. It is expected to get a memory buffer containing the
// / generated device image that could be loaded to the device directly.
Expected<std::unique_ptr<MemoryBuffer>>
run (__tgt_device_image *Image, unsigned OptLevel,
jit::PostProcessingFn PostProcessing);
private:
// / Run backend, which contains optimization and code generation.
Expected<std::unique_ptr<MemoryBuffer>> backend (Module &M, unsigned OptLevel);
// / Run optimization pipeline.
void opt (TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
unsigned OptLevel);
// / Run code generation.
void codegen (TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
raw_pwrite_stream &OS);
LLVMContext Context;
const Triple TT;
const std::string CPU;
} // namespace
// / Control environment variables.
target::StringEnvar ReplacementModuleFileName;
target::StringEnvar PreOptIRModuleFileName;
target::StringEnvar PostOptIRModuleFileName;
};
JITEngine::JITEngine (Triple::ArchType TA) : TT(Triple::getArchTypeName(TA)) {
std::call_once (InitFlag, init, TT);
}
void JITEngine::opt (TargetMachine *TM, TargetLibraryInfoImpl *TLII, Module &M,
unsigned OptLevel) {
Expand Down
Expand Up
@@ -274,18 +246,19 @@ void JITEngine::codegen(TargetMachine *TM, TargetLibraryInfoImpl *TLII,
PM.run (M);
}
Expected<std::unique_ptr<MemoryBuffer>> JITEngine::backend (Module &M,
unsigned OptLevel) {
Expected<std::unique_ptr<MemoryBuffer>>
JITEngine::backend (Module &M, const std::string &ComputeUnitKind,
unsigned OptLevel) {
auto RemarksFileOrErr = setupLLVMOptimizationRemarks (
Context , /* RemarksFilename */ " " , /* RemarksPasses */ " " ,
M. getContext () , /* RemarksFilename */ " " , /* RemarksPasses */ " " ,
/* RemarksFormat */ " " , /* RemarksWithHotness */ false );
if (Error E = RemarksFileOrErr.takeError ())
return std::move (E);
if (*RemarksFileOrErr)
(*RemarksFileOrErr)->keep ();
auto TMOrErr = createTargetMachine (M, CPU , OptLevel);
auto TMOrErr = createTargetMachine (M, ComputeUnitKind , OptLevel);
if (!TMOrErr)
return TMOrErr.takeError ();
Expand All
@@ -302,7 +275,8 @@ Expected<std::unique_ptr<MemoryBuffer>> JITEngine::backend(Module &M,
M.print (FD, nullptr );
}
opt (TM.get (), &TLII, M, OptLevel);
if (!JITSkipOpt)
opt (TM.get (), &TLII, M, OptLevel);
if (PostOptIRModuleFileName.isPresent ()) {
std::error_code EC;
Expand All
@@ -324,13 +298,26 @@ Expected<std::unique_ptr<MemoryBuffer>> JITEngine::backend(Module &M,
}
Expected<std::unique_ptr<MemoryBuffer>>
JITEngine::run (__tgt_device_image *Image, unsigned OptLevel,
jit::PostProcessingFn PostProcessing) {
JITEngine::getOrCreateObjFile (const __tgt_device_image &Image, LLVMContext &Ctx,
const std::string &ComputeUnitKind) {
// Check if the user replaces the module at runtime with a finished object.
if (ReplacementObjectFileName.isPresent ()) {
auto MBOrErr =
MemoryBuffer::getFileOrSTDIN (ReplacementObjectFileName.get ());
if (!MBOrErr)
return createStringError (MBOrErr.getError (),
" Could not read replacement obj from %s\n " ,
ReplacementModuleFileName.get ().c_str ());
return std::move (*MBOrErr);
}
Module *Mod = nullptr ;
// Check if the user replaces the module at runtime or we read it from the
// image.
// TODO: Allow the user to specify images per device (Arch + ComputeUnitKind).
if (!ReplacementModuleFileName.isPresent ()) {
auto ModOrErr = createModuleFromImage (Image, Context );
auto ModOrErr = createModuleFromImage (Image, Ctx );
if (!ModOrErr)
return ModOrErr.takeError ();
Mod = ModOrErr->release ();
Expand All
@@ -341,45 +328,79 @@ JITEngine::run(__tgt_device_image *Image, unsigned OptLevel,
return createStringError (MBOrErr.getError (),
" Could not read replacement module from %s\n " ,
ReplacementModuleFileName.get ().c_str ());
auto ModOrErr = createModuleFromMemoryBuffer (MBOrErr.get (), Context );
auto ModOrErr = createModuleFromMemoryBuffer (MBOrErr.get (), Ctx );
if (!ModOrErr)
return ModOrErr.takeError ();
Mod = ModOrErr->release ();
}
auto MBOrError = backend (*Mod, OptLevel);
if (!MBOrError)
return MBOrError.takeError ();
return backend (*Mod, ComputeUnitKind, JITOptLevel);
}
Expected<const __tgt_device_image *>
JITEngine::compile (const __tgt_device_image &Image,
const std::string &ComputeUnitKind,
PostProcessingFn PostProcessing) {
std::lock_guard<std::mutex> Lock (ComputeUnitMapMutex);
// Check if we JITed this image for the given compute unit kind before.
ComputeUnitInfo &CUI = ComputeUnitMap[ComputeUnitKind];
if (__tgt_device_image *JITedImage = CUI.TgtImageMap .lookup (&Image))
return JITedImage;
auto ObjMBOrErr = getOrCreateObjFile (Image, CUI.Context , ComputeUnitKind);
if (!ObjMBOrErr)
return ObjMBOrErr.takeError ();
return PostProcessing (std::move (*MBOrError));
auto ImageMBOrErr = PostProcessing (std::move (*ObjMBOrErr));
if (!ImageMBOrErr)
return ImageMBOrErr.takeError ();
CUI.JITImages .push_back (std::move (*ImageMBOrErr));
__tgt_device_image *&JITedImage = CUI.TgtImageMap [&Image];
JITedImage = new __tgt_device_image ();
*JITedImage = Image;
auto &ImageMB = CUI.JITImages .back ();
JITedImage->ImageStart = (void *)ImageMB->getBufferStart ();
JITedImage->ImageEnd = (void *)ImageMB->getBufferEnd ();
return JITedImage;
}
// / A map from a bitcode image start address to its corresponding triple. If the
// / image is not in the map, it is not a bitcode image.
DenseMap<void *, Triple::ArchType> BitcodeImageMap;
Expected<const __tgt_device_image *>
JITEngine::process (const __tgt_device_image &Image,
target::plugin::GenericDeviceTy &Device) {
const std::string &ComputeUnitKind = Device.getComputeUnitKind ();
// / Output images generated from LLVM backend.
SmallVector<std::unique_ptr<MemoryBuffer>, 4 > JITImages;
PostProcessingFn PostProcessing = [&Device](std::unique_ptr<MemoryBuffer> MB)
-> Expected<std::unique_ptr<MemoryBuffer>> {
return Device.doJITPostProcessing (std::move (MB));
};
// / A list of __tgt_device_image images.
std::list<__tgt_device_image> TgtImages;
} // namespace
{
std::shared_lock<std::shared_mutex> SharedLock (BitcodeImageMapMutex);
auto Itr = BitcodeImageMap.find (Image.ImageStart );
if (Itr != BitcodeImageMap.end () && Itr->second == TT.getArch ())
return compile (Image, ComputeUnitKind, PostProcessing);
}
return &Image;
}
namespace llvm {
namespace omp {
namespace jit {
bool checkBitcodeImage (__tgt_device_image *Image, Triple::ArchType TA) {
bool JITEngine::checkBitcodeImage (const __tgt_device_image &Image) {
TimeTraceScope TimeScope (" Check bitcode image" );
std::lock_guard<std::shared_mutex> Lock (BitcodeImageMapMutex);
{
auto Itr = BitcodeImageMap.find (Image-> ImageStart );
if (Itr != BitcodeImageMap.end () && Itr->second == TA )
auto Itr = BitcodeImageMap.find (Image. ImageStart );
if (Itr != BitcodeImageMap.end () && Itr->second == TT. getArch () )
return true ;
}
StringRef Data (reinterpret_cast <const char *>(Image->ImageStart ),
reinterpret_cast <char *>(Image->ImageEnd ) -
reinterpret_cast <char *>(Image->ImageStart ));
StringRef Data (reinterpret_cast <const char *>(Image.ImageStart ),
target::getPtrDiff (Image.ImageEnd , Image.ImageStart ));
std::unique_ptr<MemoryBuffer> MB = MemoryBuffer::getMemBuffer (
Data, /* BufferName */ " " , /* RequiresNullTerminator */ false );
if (!MB)
Expand All
@@ -392,37 +413,8 @@ bool checkBitcodeImage(__tgt_device_image *Image, Triple::ArchType TA) {
}
auto ActualTriple = FOrErr->TheReader .getTargetTriple ();
auto BitcodeTA = Triple (ActualTriple).getArch ();
BitcodeImageMap[Image.ImageStart ] = BitcodeTA;
if (Triple (ActualTriple).getArch () == TA) {
BitcodeImageMap[Image->ImageStart ] = TA;
return true ;
}
return false ;
return BitcodeTA == TT.getArch ();
}
Expected<__tgt_device_image *> compile (__tgt_device_image *Image,
Triple::ArchType TA, std::string MCPU,
unsigned OptLevel,
PostProcessingFn PostProcessing) {
JITEngine J (TA, MCPU);
auto ImageMBOrErr = J.run (Image, OptLevel, PostProcessing);
if (!ImageMBOrErr)
return ImageMBOrErr.takeError ();
JITImages.push_back (std::move (*ImageMBOrErr));
TgtImages.push_back (*Image);
auto &ImageMB = JITImages.back ();
auto *NewImage = &TgtImages.back ();
NewImage->ImageStart = (void *)ImageMB->getBufferStart ();
NewImage->ImageEnd = (void *)ImageMB->getBufferEnd ();
return NewImage;
}
} // namespace jit
} // namespace omp
} // namespace llvm